Machine Learning

The machine learning (ML) tracker detects tropical cyclones using a U-Net segmentation model, and then links the detections into trajectories. The model weights are downloaded from a HuggingFace Hub repository (or supplied locally) rather than wrapping an external executable.

The model and its training pipeline are described in the reference cyclone-track-ml repository.

For an overview of the functionalities, installation, and usage see the Machine Learning section of the documentation.

class tctrack.machine_learning.MLTracker(parameters, stitch_parameters=None)[source]

Bases: TCMLTracker

Class used to run the machine-learning cyclone tracking algorithm.

The model weights are loaded from a HuggingFace Hub repository on first use and cached locally by huggingface_hub.

parameters

Class containing the parameters for the tracking algorithm.

Type:

MLParameters

stitch_parameters

Class containing the parameters for the stitching step.

Type:

MLStitchParameters

model

The loaded TorchScript model in evaluation mode.

Type:

torch.jit.ScriptModule

Examples

>>> params = MLParameters(input_file="era5_2020.nc")
>>> tracker = MLTracker(params)
>>> tracker.run_tracker("trajectories.nc")
read_trajectories()[source]

Parse tracker outputs into a list of tctrack.core.Trajectory.

Returns:

A list of tctrack.core.Trajectory objects.

Return type:

list[Trajectory]

property global_metadata: dict

Read-only property containing metadata for tctrack / tracker instance.

Raises:

AttributeError – If _global_metadata has not been initialized.

Type:

dict

run_tracker_subprocess(command_name, command_list, input_file=None, input_str=None, cwd=None, verbosity=1)

Run a subprocess command for a cyclone tracking algorithm.

Parameters:

  • command_name (str) – The name of the command, used in logging and error messages.

  • command_list (list[str]) – The command and its arguments to execute.

  • input_file (str | None) – Path to a file to pass to the process via stdin (e.g. a namelist). Cannot be used together with input_str. Defaults to None.

  • input_str (str | None) – A string to pass to the process via stdin (e.g. interactive inputs). Cannot be used together with input_file. Defaults to None.

  • cwd (str | None) – Working directory in which to execute the command. Defaults to None.

  • verbosity (int) – Controls how much output is shown: 0 = No output gets printed. 1 = summary, first and last 12 lines printed (default). 2 = Entire output is streamed in real-time. Defaults to 1.

Returns:

Dictionary of subprocess output to ‘stdout’, ‘stderr’, and ‘returncode’.

Return type:

dict

Raises:

  • ValueError – If both input_file and input_str are provided simultaneously.

  • ValueError – If verbosity is not 0, 1, or 2.

set_metadata()

Initialise the metadata attributes used in the CF-netcdf output.

This is called from the to_netcdf() method.

This sets the _global_metadata attribute with the TCTrack parameters and name of the tracker. It then calls the subclass-specific _set_metadata() to set _variable_metadata and _time_metadata.

Return type:

None

property time_metadata: TCTrackerTimeMetadata

Read-only property containing time metadata for this Tracker run.

Raises:

AttributeError – If _time_metadata has not been initialized.

Type:

dict

to_netcdf(output_file, trajectories=None)

Write track trajectories to CF-compliant NetCDF trajectory file.

Reads in trajectories based on the parameters set for a specific implementation of the class and writes them to a CF-Conventions compliant NetCDF trajectory file using cf-python. Trajectories are assumed to contain as a minimum data for lat, lon, and timestep.

An ancillary field variable is added to the output file indicating any tracks that start/end within 1 day of the input dataset boundaries.

Parameters:

  • output_file (str) – filename for the output netCDF file Note: This will be placed in the local directory unless a full path is given

  • trajectories (list[Trajectory] | None) – An optional list of trajectories to write. If the tracking algorithm produces an intermediate output file then this can be skipped as it will instead read from there.

Return type:

None

Warning

UserWarning

If there are no trajectories read in from the tracker outputs by read_trajectories() meaning no NetCDF output file can be written.

References

CF-Conventions v1.1 - H.4. Trajectory Data cf-python documentation

Examples

Instantiate a TCTracker subclass instance with appropriate parameters, run the relevant methods to generate cyclone track trajectories, and then save the results to a CF-compliant trajectory file:

>>> my_tracker = TCTracker(...)
>>> ...
>>> my_tracker.to_netcdf("my_netcdf_file.nc")
property variable_metadata: dict

Read-only property containing NetCDF metadata for variables.

Raises:

AttributeError – If _variable_metadata has not been initialized.

Type:

dict

preprocess()[source]

Build the model input tensor from the configured ERA5 variables.

Reproduces the channel layout the model was trained on: each of parameters.pressure_variables at each of parameters.pressure_levels, followed by a land/sea mask, followed by sea surface temperature gap-filled with 2-metre temperature over land. Channels are then normalised using _load_normalisation_stats.

Returns:

Tensor of shape (channel, time, lat, lon), normalised.

Return type:

torch.Tensor

Raises:

ValueError – If a configured variable cannot be found in the input file.

detect()[source]

Run the U-Net over the grid and locate the tropical cyclones in it.

For each timestep the model is run, its per-pixel class probabilities are thresholded, and the surviving pixels are grouped into discrete candidate storms (see _cluster_candidates() and _merge_nearby_candidates()). The result is a list of TC locations on _candidates, which stitch() then links into trajectories - mirroring how the other TCTrack trackers split the work, where detection emits candidate points and stitching consumes them.

The full per-pixel probabilities are also kept on _scores, which retains the per-class detail that thresholding discards (useful for evaluation, e.g. one-vs-rest ROC per lifecycle class).

Raises:

ValueError – If the model’s output spatial shape doesn’t match the input file’s latitude/longitude grid, it would indicate a bug.

Return type:

None

stitch()[source]

Link the TC locations found by detect() into trajectories.

Consumes the candidates on _candidates, working through the timesteps in order and, at each one, extending existing tracks with a nearby candidate, retiring tracks that have gone unmatched for too long, and starting new tracks from whatever is left over.

The result is also stored on _trajectories, so that read_trajectories() - and therefore to_netcdf() - can reach it.

Returns:

One Trajectory per track with at least stitch_parameters.min_length observations.

Return type:

list[Trajectory]

detections_to_netcdf(output_file)[source]

Write the detected TC locations to a CF-compliant NetCDF point file.

The detections from detect() are a set of independent points - each a storm location at one time, with the input variables sampled there - rather than connected tracks. They are therefore written with featureType = "point" over a single observation dimension, instead of the trajectory layout to_netcdf() produces.

Parameters:

output_file (str) – Filename for the output NetCDF file. Placed in the working directory unless a full path is given.

Warns:

UserWarning – If detect() found no candidates, in which case no file is written.

Return type:

None

References

CF-Conventions - H.1. Point Data

run_tracker(output_file)[source]

Run the tracker to obtain the tropical cyclone trajectories.

Runs detect() to locate the cyclones at each timestep - which reads and prepares the input itself, via preprocess() - then links those detections into tracks with stitch(), and finally writes the tracks out with to_netcdf().

The detections behind those tracks can be saved separately, as CF point data, with detections_to_netcdf().

Parameters:

output_file (str) – Filename to which the tropical cyclone track trajectories are saved.

Return type:

None

Examples

To set the parameters, instantiate a MLTracker instance and run the algorithm to generate trajectories:

>>> params = MLParameters(...)
>>> tracker = MLTracker(params)
>>> tracker.run_tracker("trajectories.nc")
class tctrack.machine_learning.MLParameters(model_path=None, hf_repo_id='surbhigoel456/cyclone-TC-ML', hf_token=None, device='cpu', threshold=0.5, input_file='', pressure_variables=('relative_humidity', 'air_temperature', 'eastward_wind', 'northward_wind', 'atmosphere_relative_vorticity'), pressure_levels=(1000, 750, 500), sst_variable='ncvar%sst', t2m_variable='ncvar%t2m', normalisation_stats_path=None, merge_distance_deg=2.0)[source]

Bases: TCMLParameters

Dataclass containing values for parameters used by MLTracker.

Extends TCMLParameters with the input file path and sets the default HuggingFace repository for this model.

See also

MLTracker

The tracker class that uses these parameters.

input_file: str = ''

CF-NetCDF file containing the input variables for the model.

hf_repo_id: str = 'surbhigoel456/cyclone-TC-ML'

HuggingFace Hub repository ID for the cyclone detection model.

pressure_variables: Sequence[str] = ('relative_humidity', 'air_temperature', 'eastward_wind', 'northward_wind', 'atmosphere_relative_vorticity')

CF identities of the pressure-level input variables, in model channel order.

Correspond to pressure_var_codes in the reference training pipeline (cyclone-track-ml), but given as the actual CF standard_name each variable is converted to (via cfgrib) in ERA5 output, which differs from the CDS API request name used there for several variables (e.g. "temperature" in the reference pipeline is "air_temperature" here).

pressure_levels: Sequence[float] = (1000, 750, 500)

Pressure levels (hPa) to extract for each of pressure_variables.

sst_variable: str = 'ncvar%sst'

CF identity of the sea surface temperature variable.

ERA5 sea surface temperature has no standard_name (cfgrib leaves it "unknown"), so it must be selected by NetCDF variable name instead.

t2m_variable: str = 'ncvar%t2m'

CF identity of the 2-metre temperature, used to gap-fill SST over land.

normalisation_stats_path: str | None = None

Path to a .nc file with per-channel mean and range variables (each shape (channel,), ordered as in MLTracker.preprocess()), used to normalise input as (x - mean) / range.

These must be the statistics computed over the model’s training set; they cannot be recomputed from a single inference file. Matches the data/normalisation_parameters.nc file saved by compiler.py in the reference training pipeline. If None, the statistics bundled with this package - computed over the default model’s training set - are used instead.

merge_distance_deg: float = 2.0

Distance (degrees lat/lon) within which candidates in the same timestep are merged, keeping only the highest-scoring one.

The model often splits a single storm into several adjacent clusters - either side of a lifecycle-class boundary, or where confidence dips below threshold mid-storm - which would otherwise each become a separate track. Mirrors merge_dist in TempestExtremes’ DetectNodes, which resolves the same problem for its own candidates. Set to 0 to disable merging.

device: str = 'cpu'

PyTorch device to run inference (e.g. "cpu" or "cuda").

hf_token: str | None = None

HuggingFace access token for the model repository. If None, the value of the HF_TOKEN environment variable is used instead.

model_path: str | None = None

Local path to the .pt model file. If None, the model is downloaded from the HuggingFace Hub repository given by hf_repo_id.

threshold: float = 0.5

Confidence threshold for retaining detections. Candidates with a model score below this value are discarded.

class tctrack.machine_learning.MLStitchParameters(max_distance_deg=3.0, max_gap=1, min_length=2)[source]

Bases: TCTrackerParameters

Dataclass containing values used by MLTracker.stitch().

See also

MLTracker

The tracker class that uses these parameters.

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.

max_gap: int = 1

Number of consecutive timesteps a track may go unmatched before it is closed, tolerating brief drops below threshold.

min_length: int = 2

Minimum number of observations a track must have to be kept, filtering out single-timestep detections likely to be noise.

The classes below are base classes for adding other machine learning tracking algorithms to TCTrack. They are not intended to be used directly.

class tctrack.core.ml_tracker.TCMLParameters(model_path=None, hf_repo_id='', hf_token=None, device='cpu', threshold=0.5)[source]

Bases: TCTrackerParameters

Base dataclass for parameters common to all ML-based trackers.

model_path: str | None = None

Local path to the .pt model file. If None, the model is downloaded from the HuggingFace Hub repository given by hf_repo_id.

hf_repo_id: str = ''

HuggingFace Hub repository ID from which the model is downloaded.

hf_token: str | None = None

HuggingFace access token for the model repository. If None, the value of the HF_TOKEN environment variable is used instead.

device: str = 'cpu'

PyTorch device to run inference (e.g. "cpu" or "cuda").

threshold: float = 0.5

Confidence threshold for retaining detections. Candidates with a model score below this value are discarded.

class tctrack.core.ml_tracker.TCMLTracker[source]

Bases: TCTracker

Abstract base class for ML-based tropical cyclone trackers.

Extends TCTracker functionality common to all ML-based trackers e.g. loading a TorchScript model from HuggingFace Hub repository, and the detect step.

model

The loaded TorchScript model in evaluation mode.

Type:

torch.jit.ScriptModule

abstract detect()[source]

Run the ML model on the preprocessed input to obtain TC candidates.

This method should populate the internal state (e.g. a list of candidate detections) that can be used for reading trajectories or the stitching step.

Return type:

None