Tracker module for ML based TC tracking algortihms by surbhigoel77 · Pull Request #207 · Cambridge-ICCS/TCTrack

surbhigoel77 · 2026-06-02T17:50:29Z

Linked issue: #190

The PR aims to build the core infrastructure for running ML models within the TCTrack framework. It uses exisiting base classes and defines new ml-specific classes.

It includes the following files (based on #204):

ml_tracker.py : contains a base dataclass TCMLParameters to fetch ML Parameters and a base abstract class TCMLTracker that extends the existing Tracker by adding a load_model function (for ml models). These classes remain common for all ML models (basically it is a way to standardise processes for any incoming new model).
cyclone_track_ml.py : inherits the base classes defined in ml_tracker.py and adds new model specific functions. Reads the model from hugging face hub repository and connects it to the run_tracker pipeline. For classes that are common for both ML and non-ML models like TCTrackerMetadata and Trajectory are inherited from the original core/tracker.py.
'run_tracker_on_real_data.py' : script to run the ML Tracker on a subset of ERA5 data.
test_ml_tracker.py : unit, integration (Still in progress)

Assumptions:

merge_distance_deg = 2.0 (How close (in degrees) two candidates in the same timestep can be before the weaker one is discarded as a duplicate)
stitch_max_distance_deg = 3.0 (How far (in degrees) a candidate can be from a track's last known position and still be claimed as that same storm continuing)
stitch_max_gap: int = 1 (How many timesteps in a row a track can go unmatched before it's retired)
stitch_min_length: int = 2 (The minimum number of points a track must have to be kept in the final output)

…rence

…acks

surbhigoel77 · 2026-08-28T08:35:19Z

@sjavis you can try and run the ML Tracker using the 'run_tracker_on_real_data.py' script. You will need ERA5 data and the normalisation stats to run this. I have already provided the absolute paths/creds in the code for

the hugging face repo for model weights, and
the required ERA5 data and norm stats

Let me know if have issues accessing it.

The integration test is not to be reviewed atm, I am making changes to it

sjavis

I haven't reviewed the stitching or the integration test yet but I thought I would give my suggestions so far. Overall the structure seems good.

sjavis · 2026-09-01T14:43:14Z

+    stitch_max_distance_deg: float = 3.0
+    """Maximum distance (degrees lat/lon) between a track's last point and a
+    candidate in the next timestep for them to be linked as the same storm.
+
+    Not part of the reference training pipeline, which only trains
+    per-pixel classification and has no timestep-to-timestep linking step.
+    The default assumes 6-hourly input (as used for training) and typical
+    cyclone translation speeds; reduce for higher-frequency input.
+    """
+
+    stitch_max_gap: int = 1
+    """Number of consecutive timesteps a track may go unmatched before it is
+    closed, tolerating brief drops below :attr:`~TCMLParameters.threshold`.
+    """
+
+    stitch_min_length: int = 2
+    """Minimum number of observations a track must have to be kept, filtering
+    out single-timestep detections likely to be noise.
+    """


I suggest making a separate class for the stitching parameters.

This will also need adding to the __init__.py file

And the parameters can probably loose the stitch_ prefix since they are now in the class.

sjavis · 2026-09-04T08:09:05Z

+        self._trajectories: list[Trajectory] = []
+        self._scores: list[dict] = []
+        # TC locations found by detect(), consumed by stitch().
+        self._candidates: list[dict] = []


I suggest we make a candidate class (maybe using a TypedDict) since it is used in a few places and it would make the typing clearer than just using a dict.

This is looking good and I realise it was more complicated than I anticipated as I thought extra_keys was available (rather than being new in python 3.15).

I suggest making the extra variables inside a data dictionary to resolve the need for casting, etc. I will make a separate PR to show you my suggestion as it requires changes in a few places. (Added in #259)

sjavis · 2026-09-04T08:55:55Z

Just to clarify, is this file intended just for development and to be removed before merging?

I want to add it to the tutorials as we have tutorial scripts for all other models, so might as well add a script for the ml one.

sjavis

This is looking good. Thank you for the changes that you previously made. I've left some comments and will open a separate PR for the changes to Candidate that I suggested.

sjavis · 2026-09-15T14:42:39Z

+    stitch_max_distance_deg: float = 3.0
+    """Maximum distance (degrees lat/lon) between a track's last point and a
+    candidate in the next timestep for them to be linked as the same storm.
+
+    Not part of the reference training pipeline, which only trains
+    per-pixel classification and has no timestep-to-timestep linking step.
+    The default assumes 6-hourly input (as used for training) and typical
+    cyclone translation speeds; reduce for higher-frequency input.
+    """
+
+    stitch_max_gap: int = 1
+    """Number of consecutive timesteps a track may go unmatched before it is
+    closed, tolerating brief drops below :attr:`~TCMLParameters.threshold`.
+    """
+
+    stitch_min_length: int = 2
+    """Minimum number of observations a track must have to be kept, filtering
+    out single-timestep detections likely to be noise.
+    """


This will also need adding to the __init__.py file

sjavis · 2026-09-15T16:32:48Z

+    Plain ``hypot(dlat, dlon)`` on raw coordinate differences is wrong in two
+    ways: it overstates distances that cross the antimeridian (e.g. 179E to
+    179W is 2 degrees apart, not 358), and it ignores that a degree of
+    longitude covers less physical distance away from the equator. This
+    wraps the longitude difference to ``(-180, 180]`` and scales it by the
+    cosine of the mean latitude to correct for both. Still an approximation
+    - the Haversine formula would be exact - but adequate at the latitudes
+    tropical cyclones occur at.


Suggested change

Plain ``hypot(dlat, dlon)`` on raw coordinate differences is wrong in two

ways: it overstates distances that cross the antimeridian (e.g. 179E to

179W is 2 degrees apart, not 358), and it ignores that a degree of

longitude covers less physical distance away from the equator. This

wraps the longitude difference to ``(-180, 180]`` and scales it by the

cosine of the mean latitude to correct for both. Still an approximation

- the Haversine formula would be exact - but adequate at the latitudes

tropical cyclones occur at.

An approximation - the Haversine formula would be exact - but adequate at small

distances and the latitudes tropical cyclones occur at.

A smaller description is probably sufficient.

sjavis · 2026-09-16T11:14:21Z

+        self._trajectories: list[Trajectory] = []
+        self._scores: list[dict] = []
+        # TC locations found by detect(), consumed by stitch().
+        self._candidates: list[dict] = []


This is looking good and I realise it was more complicated than I anticipated as I thought extra_keys was available (rather than being new in python 3.15).

I suggest making the extra variables inside a data dictionary to resolve the need for casting, etc. I will make a separate PR to show you my suggestion as it requires changes in a few places. (Added in #259)

sjavis · 2026-09-16T11:25:00Z

+import torch
+from huggingface_hub import hf_hub_download


torch and huggingface_hub should be added as dependencies to the pyproject.toml

sjavis · 2026-09-16T12:59:03Z

+    MLTracker,
+)
+
+RDS_DIR = "/home/sg2147/rds/rds-inspire-tc-TqEGHMWTn8A/sg2147"


Suggested change

RDS_DIR = "/home/sg2147/rds/rds-inspire-tc-TqEGHMWTn8A/sg2147"

RDS_DIR = "/rds/project/rds-TqEGHMWTn8A/sg2147"

Longer term it would be best if we could make this not need data on CSD3

sjavis · 2026-09-16T13:05:35Z

+            model_path=MODEL_PATH,
+            normalisation_stats_path=stats_file,


Can these arguments be ignored here or are the files different to the defaults?

sjavis · 2026-09-16T14:02:43Z

+def fetch_norm_stats(stats_file: str) -> None:
+    """Pull the real normalisation statistics from the reference repo."""
+    git = ["git", "-C", REFERENCE_REPO]
+    subprocess.run([*git, "fetch", "origin"], check=True)  # noqa: S603
+    with open(stats_file, "wb") as handle:
+        subprocess.run(  # noqa: S603
+            [*git, "show", REFERENCE_STATS_REF], stdout=handle, check=True
+        )
+    print(f"  fetched statistics from {REFERENCE_STATS_REF}")


Should this now be able to use the bundled stats file?

sjavis · 2026-09-16T14:13:21Z

+    print(f"  fetched statistics from {REFERENCE_STATS_REF}")
+
+
+def test_pipeline(input_file: str, stats_file: str) -> None:


I would suggest converting this and test_output_writers to be compatible with pytest. But make it skip them by default with something like:

@pytest.mark.skip(reason="Requires data on CSD3")

sjavis · 2026-09-16T14:24:20Z

+    trajectories = tracker.stitch()
+    assert tracker.read_trajectories() == trajectories, (
+        "read_trajectories() must return what stitch() produced, else to_netcdf() "
+        "would write nothing"
+    )
+    for trajectory in trajectories:
+        assert trajectory.observations >= tracker.stitch_parameters.stitch_min_length
+        for key in ("time", "lat", "lon"):
+            assert key in trajectory.data, f"trajectory missing '{key}'"
+    print(f"  stitch():     {len(trajectories)} trajectories, all well-formed   OK")
+    if not trajectories:
+        print("                (none found - fine, model quality is not under test)")


This should probably check that there is at least one trajectory detected (although I don't know if any would be detected from the input data). It doesn't really check anything otherwise.

Add placeholder file for ml-tracker

1d31c7b

surbhigoel77 marked this pull request as draft June 2, 2026 17:50

This was referenced Jun 2, 2026

Meta: ML detection algorithm #190

Open

ML Tracker skeleton #195

Closed

surbhigoel77 self-assigned this Jun 3, 2026

surbhigoel77 added 7 commits June 17, 2026 15:35

Add parameter class for the model, including hugging face details

496c768

Add ML_Tracker class init function

7210a2d

Add parameter class for ML trackers

13ae971

Add tracker class for ml models, model load function

1dbc27e

Add detect function to ML tracker class

a6ca45a

Update imports: use TCMLParameters and TCMLTracker base classes

909cc47

Update MLParameters: inherit common fields from TCMLParameters class

68bcc19

surbhigoel77 added enhancement New feature or request MLTracker labels Jun 24, 2026

Refactor MLTracker

82ef996

This was referenced Jul 8, 2026

Add ML detection neural net #196

Open

Add preprocessing for ML tracker #197

Open

surbhigoel77 added 13 commits July 24, 2026 11:14

Add UNet parameters to MLparameters class

fef89eb

Add preprocessing skeleton code

e00333c

Add fucntions: load norm data, preprocessing of inpur tensor for infe…

f05e2bd

…rence

Add detect step structure inspired from cyclone genesis model

51c94bf

Update load_normalisation function to enable reading .nc norm files

16b3382

Import scipy before torch to prevent GLIBCXX import crash

f58d9d3

Add fucntions to cluster candidates and merge the neigbours

f766557

Add fucntion to find closest unmatched candidate

157c5f7

Add stitch fucntion that uses helper fucntion to build possible TC tr…

5a42284

…acks

Add integration test for cyclone-genesis ml model

eea5323

Move candidate detection from stitch into detect

a99f6ef

Record co-located variables at each detected TC location

b6c28a0

Add detections_to_netcdf to save TC locations as CF point data

4f19df1

surbhigoel77 added 2 commits August 11, 2026 21:12

Include run tracker in test

13597e9

Add helper comments

0b05a26

This was linked to issues Aug 18, 2026

Add ML detection neural net #196

Open

Add preprocessing for ML tracker #197

Open

Locate the storm centres in the ML detection #200

Open

Save detection outputs in ML tracker #198

Open

Add a basic stitching function #201

Open

Add a script to run ML Tracker with a subset of data

b7df152

sjavis mentioned this pull request Aug 27, 2026

Adds template for new tracking algorithms #204

Merged

3 tasks done

surbhigoel77 requested a review from sjavis August 28, 2026 08:30

sjavis reviewed Sep 4, 2026

View reviewed changes

surbhigoel77 added 11 commits September 14, 2026 11:01

Add init for ML Tracker

8692f64

Add init for ML Trscker module

0467d39

Split stitching parameters into MLStitchParameters

bd86814

Update integration test for MLStitchParameters split

8518d4e

Clear HF token from parameters before it reaches output metadata

aad284f

Bundle default normalisation stats with the ML tracker package

0a61b45

Add Candidate TypedDict for detection dicts

c9ed406

Simplify unvisited-pixel set construction in _cluster_candidates

29f7f2e

Cluster storm pixels regardless of class, take class from peak pixel

d868bba

Fix antimeridian and latitude-scaling bug in candidate distance calc

c3c8de1

Clarify sort-priority comment in _merge_nearby_candidates

4c4c6bb

surbhigoel77 requested a review from sjavis September 15, 2026 09:24

sjavis requested changes Sep 16, 2026

View reviewed changes

surbhigoel77 added 5 commits September 29, 2026 11:38

Export MLStitchParameters from tctrack.machine_learning

52f4a25

Drop redundant stitch_ prefix from MLStitchParameters fields

d590dd4

Simplify coordinate long_names in detections output

2c4e43d

Clarify approximation note in _angular_distance_deg docstring

bb426ab

Add hf,scipy,torch as dependencies

37bd666

	RDS_DIR = "/home/sg2147/rds/rds-inspire-tc-TqEGHMWTn8A/sg2147"
	RDS_DIR = "/rds/project/rds-TqEGHMWTn8A/sg2147"

		print(f" fetched statistics from {REFERENCE_STATS_REF}")


		def test_pipeline(input_file: str, stats_file: str) -> None:

Conversation

surbhigoel77 commented Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

surbhigoel77 commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sjavis left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sjavis Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sjavis left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sjavis Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

surbhigoel77 commented Jun 2, 2026 •

edited

Loading

surbhigoel77 commented Aug 28, 2026 •

edited

Loading

sjavis Sep 16, 2026 •

edited

Loading

sjavis Sep 16, 2026 •

edited

Loading