"""Module providing classes that bind the TSTORMS code.
References
----------
- TSTORMS NOAA page: https://www.gfdl.noaa.gov/tstorms/
"""
import os
import tempfile
import textwrap
import warnings
from typing import Iterable
import cf
from cftime import num2date
from netCDF4 import Dataset
from tctrack.core import TCTracker, TCTrackerMetadata, TCTrackerParameters, Trajectory
from tctrack.tstorms import (
TSTORMSBaseParameters,
TSTORMSDetectParameters,
TSTORMSStitchParameters,
)
[docs]
class TSTORMSTracker(TCTracker):
"""Class containing bindings to the TSTORMS code.
Attributes
----------
driver_parameters : TSTORMSDetectParameters | None
Class containing the parameters for the driver detection algorithm
trajectory_parameters : TSTORMSStitchParameters | None
Class containing the parameters for the trajectory stitching algorithm
"""
# Private attributes
_tempdir: tempfile.TemporaryDirectory
_u_in_file: str | None = None
_v_in_file: str | None = None
_vort_in_file: str | None = None
_tm_in_file: str | None = None
_slp_in_file: str | None = None
def __init__(
self,
tstorms_parameters: TSTORMSBaseParameters,
detect_parameters: TSTORMSDetectParameters | None = None,
stitch_parameters: TSTORMSStitchParameters | None = None,
):
"""
Construct the TSTORMSTracker class.
Parameters
----------
tstorms_parameters : TSTORMSBaseParameters
Class containing the parameters for setting up TSTORMS usage.
detect_parameters : TSTORMSDetectParameters | None
Class containing the parameters for the node detection in TSTORMS.
Used to create the namelists for `tstorms_driver`.
Defaults to the default values in TSTORMSDetectParameters Class
stitch_parameters : TSTORMSStitchParameters | None
Class containing the parameters for the stitching algorithm of TSTORMS
Defaults to the default values in TSTORMSStitchParameters Class
"""
self.tstorms_parameters: TSTORMSBaseParameters = tstorms_parameters
if detect_parameters is not None:
self.detect_parameters: TSTORMSDetectParameters = detect_parameters
else:
self.detect_parameters = TSTORMSDetectParameters()
if stitch_parameters is not None:
self.stitch_parameters: TSTORMSStitchParameters = stitch_parameters
else:
self.stitch_parameters = TSTORMSStitchParameters()
# Set any manually defined input file attributes
for var in ["u", "v", "vort", "tm", "slp"]:
file_param = getattr(self.detect_parameters, f"{var}_in_file")
if file_param is not None:
setattr(
self,
f"_{var}_in_file",
os.path.join(self.tstorms_parameters.input_dir, file_param),
)
# Ensure the output directory exists, create if not.
output_dir = self.tstorms_parameters.output_dir
os.makedirs(output_dir, exist_ok=True)
@property
def _parameters(self) -> list[TCTrackerParameters]:
"""A list of the parameter objects that is accessible from the base class."""
return [self.tstorms_parameters, self.detect_parameters, self.stitch_parameters]
def _in_file(self, var: str) -> str:
"""Get the filepath for one of the input files."""
file_attr = f"_{var}_in_file"
if getattr(self, file_attr) is None:
msg = (
f"`{file_attr}` not set. Ensure that the input files include all "
"required variables and that `set_input_files` has been called."
)
raise RuntimeError(msg)
return getattr(self, file_attr)
def _write_driver_namelist(self) -> str:
"""
Generate the namelist file for the tstorms_driver routine.
The namelist file will be written to the ``output_dir`` specified in
:attr:`tstorms_parameters`.
Returns
-------
str
The full path to the generated namelist file.
Raises
------
FileNotFoundError
If the `tstorms_dir` does not exist.
"""
# Ensure the output_dir exists to place namelist in
output_dir = self.tstorms_parameters.output_dir
if not os.path.exists(output_dir):
err_msg = f"TSTORMS output directory '{output_dir}' does not exist."
raise FileNotFoundError(err_msg)
namelist_path = os.path.join(output_dir, "nml_driver")
# Format the namelist content
detect_params = self.detect_parameters
namelist_content = textwrap.dedent(f"""
&nml_tstorms
crit_vort = {detect_params.vort_crit:.4E}
crit_twc = {detect_params.tm_crit:.4f}
crit_thick = {detect_params.thick_crit:.4f}
crit_dist = {detect_params.dist_crit:.4f}
lat_bound_n = {detect_params.lat_bound_n:.4f}
lat_bound_s = {detect_params.lat_bound_s:.4f}
do_spline = {".true." if detect_params.do_spline else ".false."}
do_thickness= {".true." if detect_params.do_thickness else ".false."}
&end
&input
fn_u = '{self._in_file("u")}'
fn_v = '{self._in_file("v")}'
fn_vort = '{self._in_file("vort")}'
fn_tm = '{self._in_file("tm")}'
fn_slp = '{self._in_file("slp")}'
use_sfc_wnd = {".true." if detect_params.use_sfc_wind else ".false."}
&end
""")
# Write the namelist content to the file
with open(namelist_path, "w") as namelist_file:
namelist_file.write(namelist_content)
return namelist_path
def _make_driver_call(self):
"""
Construct a driver call based on options set in driver_parameters.
Makes a call to write the namelist and then construct and call the driver
command.
Returns
-------
list[str]
list of strings that can be combined to form a driver command
based on the parameters set in self.detect_parameters
"""
tstorms_driver_dir = os.path.join(
self.tstorms_parameters.tstorms_dir, "tstorms_driver"
)
dn_argslist = [os.path.join(tstorms_driver_dir, "tstorms_driver.exe")]
return dn_argslist
[docs]
def detect(self, verbosity: int = 1):
"""
Call the driver utility of TSTORMS.
This will make a system call out to the tstorms_driver code from TSTORMS
(provided it has been installed as an external dependency). This will be
run according to the parameters in the :attr:`detect_parameters` attribute
that were set when the :class:`TSTORMSTracker` instance was created.
The output file is a plain text file named ``cyclones`` containing each of the
TC candidates at each time from the input files. This will be generated
in the desired output location provided in the base parameters attribute
:attr:`tstorms_parameters`.
Cyclones in the file are listed for each time in the format:
.. code-block:: text
<day> <month> <year> <number> <hour>
<i_index> <j_index> <lon_slpmin> <lat_slpmin>
<max_wind> <max_vort> <min_slp> <exist_warm_core> <exist_thickness>
<max_warm_core> <max_thickness>
...
- ``number`` is the number of nodes at that time.
- ``i_index``, ``j_index`` are the grid indices of the node.
- ``lon_slpmin``, ``lat_slpmin`` are the coordinates of the slp minimum.
- Other values are the variable values for the cyclone output from the code.
Parameters
----------
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
-------
dict
dict of subprocess output corresponding to stdout, stderr, and returncode.
Raises
------
FileNotFoundError
If the tstorms_driver executeable from TSTORMS cannot be found.
RuntimeError
If the tstorms_driver executeable from TSTORMS returns a non-zero exit code.
Examples
--------
To set the parameters, instantiate a :class:`TSTORMSTracker` instance and run
detect:
>>> base_params = TSTORMSBaseParameters(...)
>>> detect_params = TSTORMSDetectParameters(...)
>>> my_tracker = TSTORMSTracker(
tstorms_parameters=base_params,
detect_parameters=detect_params
)
>>> result = my_tracker.detect()
"""
namelist_filepath = self._write_driver_namelist()
driver_call_list = self._make_driver_call()
process_output = self.run_tracker_subprocess(
"Detect",
driver_call_list,
namelist_filepath,
verbosity=verbosity,
cwd=self.tstorms_parameters.output_dir,
)
return process_output
def _write_trajectory_analysis_namelist(self) -> str:
"""Generate the namelist file for the trajectory_analysis routine.
The namelist file will be written to the ``output_dir`` specified in
:attr:`tstorms_parameters`.
Returns
-------
str
The full path to the generated namelist file.
Raises
------
FileNotFoundError
If the `trajectory_analysis` directory does not exist.
"""
tstorms_dir = self.tstorms_parameters.tstorms_dir
# Ensure the output_dir exists to place namelist in
output_dir = self.tstorms_parameters.output_dir
if not os.path.exists(output_dir):
err_msg = f"TSTORMS output directory '{output_dir}' does not exist."
raise FileNotFoundError(err_msg)
# Define the namelist file path
namelist_path = os.path.join(output_dir, "nml_traj")
# Format the namelist content
stitch_params = self.stitch_parameters
namelist_content = textwrap.dedent(f"""
&input
rcrit = {stitch_params.r_crit:.4f}
wcrit = {stitch_params.wind_crit:.4f}
vcrit = {stitch_params.vort_crit:.4E}
twc_crit = {stitch_params.tm_crit:.4f}
thick_crit = {stitch_params.thick_crit:.4f}
nwcrit = {stitch_params.n_day_crit:.4f}
do_filt = {".true." if stitch_params.do_filter else ".false."}
nlat = {stitch_params.lat_bound_n:.4f}
slat = {stitch_params.lat_bound_s:.4f}
do_spline = {".true." if stitch_params.do_spline else ".false."}
do_thickness = {".true." if stitch_params.do_thickness else ".false."}
landmask = '{os.path.join(tstorms_dir, "trajectory_analysis/landsea.map")}'
cmask = '{os.path.join(tstorms_dir, "trajectory_analysis/imask_2")}'
&end
""")
# Write the namelist content to the file
with open(namelist_path, "w") as namelist_file:
namelist_file.write(namelist_content)
return namelist_path
def _make_trajectory_analysis_call(self):
"""
Construct a trajectory_analysis call based on options set in stitch_parameters.
Makes a call to write the namelist and then construct and call the
trajectory_analysis command.
Returns
-------
list[str]
list of strings that can be combined to form a trajectory_analysis command
based on the parameters set in self.stitch_parameters
"""
tstorms_trajectory_analysis_dir = os.path.join(
self.tstorms_parameters.tstorms_dir, "trajectory_analysis"
)
stitch_argslist = [
os.path.join(tstorms_trajectory_analysis_dir, "trajectory_analysis_csc.exe")
]
return stitch_argslist
[docs]
def stitch(self, verbosity: int = 1):
"""
Call the trajectory analysis utility of TSTORMS to stitch candidate storms.
This will make a system call out to the trajectory_analysis code from TSTORMS
(provided it has been installed as an external dependency). This will be
run according to the parameters in the :attr:`stitch_parameters` attribute
that were set when the :class:`TSTORMSTracker` instance was created.
It assumes that the candidate storms are contained in the output_dir
directory in a file ``cyclones`` text file.
The outputs are plain text files named:
- ``ori`` containing the origins of each storm in the form
``lon, lat, YY, MM, DD, HH``,
- ``traj`` containing trajectory point data in the form
``lon, lat, wind, psl, YY, MM, DD, HH`` for each trajectory,
- ``trav`` containing trajectory point data and vorticity in the form
``lon, lat, wind, psl, vort_max, YY, MM, DD, HH`` for each trajectory,
- ``stats`` containing information about the number of storms per month/year
in each basin.
If filtering is applied (:attr:`stitch_parameters` ``do_filter``) there will be
additional files ``ori_filt``, ``traj_filt``, and ``trav_filt`` from after this
takes place.
These files will be generated in the desired output location provided in the
base parameters attribute :attr:`tstorms_parameters`.
Parameters
----------
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
-------
dict
dict of subprocess output corresponding to stdout, stderr, and returncode.
Raises
------
FileNotFoundError
If the trajectory_analysis executeable from TSTORMS cannot be found.
RuntimeError
If the trajectory_analysis executeable from TSTORMS returns a non-zero
exit code.
Examples
--------
To set the parameters, instantiate a :class:`TSTORMSTracker` instance and run
stitch:
>>> base_params = TSTORMSBaseParameters(...)
>>> stitch_params = TSTORMSStitchParameters(...)
>>> stitch_params = TSTORMSTracker(
tstorms_parameters=base_params,
stitch_parameters=stitch_params
)
>>> result = my_tracker.stitch()
"""
# Check if the cyclones file exists before proceeding
cyclones_file = os.path.join(self.tstorms_parameters.output_dir, "cyclones")
if not os.path.exists(cyclones_file):
err_msg = (
"No cyclones file found in the output directory. "
"Did you forget to run `detect` or delete the file?"
)
raise FileNotFoundError(err_msg)
namelist_filepath = self._write_trajectory_analysis_namelist()
trajectory_call_list = self._make_trajectory_analysis_call()
process_output = self.run_tracker_subprocess(
"Stitch",
trajectory_call_list,
namelist_filepath,
verbosity=verbosity,
cwd=self.tstorms_parameters.output_dir,
)
return process_output
def _set_metadata(self) -> None:
"""
Set the time and variable metadata attributes by reading from input files.
Reads metadata for each variable from the input NetCDF files
defined in :attr:`detect_parameters`.
These will be stored in the :attr:`variable_metadata` attribute as a
dictionary of :class:`TCTrackerMetadata` objects.
This will be called from the :meth:`set_metadata` method.
Raises
------
ValueError
If a variable is not found in the input files.
Examples
--------
To generate in the metadata for variables from parameters and inputs:
>>> tracker = TSTORMSTracker(tstorms_params, detect_params, stitch_params)
>>> tracker.set_metadata()
>>> tracker.variable_metadata
{
"wind_speed": TCTrackerMetadata(
properties={
"standard_name": "wind_speed",
"long_name": "Wind Speed",
...
},
)
...
}
"""
if not self._time_metadata:
self._set_time_metadata()
if not self._variable_metadata:
self._set_variable_metadata()
def _set_variable_metadata(self) -> None:
"""
Extract variable metadata from inputs and return dict for setting.
TSTORMS calculates windspeed as Euclidian norm of u and v, so read
key units and attributes from u and assume they apply.
Vorticity, Temperature, and Sea-level-pressure all read in directly from file.
Note that there is some slight misdirection here in that wind, slp, and tm are
detected relative to the vorticity location, but the slp location is used to
provide coordinates. This means a situation could arise in which points are
in fact from an area more than quoted (up to twice the radius).
Raises
------
FileNotFoundError
If an input file cannot be found.
ValueError
If a variable is not found in the input files.
"""
detect_params = self.detect_parameters
var_outputs: list[dict] = [
{
"std_name": "wind_speed",
"long_name": (
"Surface Wind Speed" if detect_params.use_sfc_wind else "Wind Speed"
),
"tstorms_name": ("u_ref" if detect_params.use_sfc_wind else "u850"),
"filename": self._in_file("u"),
"cellmethod": cf.CellMethod(
"area",
"maximum",
qualifiers={
"comment": (
f"lesser circle of radius {detect_params.dist_crit} degrees"
),
},
),
},
{
"std_name": "air_pressure_at_mean_sea_level",
"long_name": "Air Pressure at Mean Sea Level",
"tstorms_name": "slp",
"filename": self._in_file("slp"),
"cellmethod": cf.CellMethod(
"area",
"maximum",
qualifiers={
"comment": (
f"lesser circle of radius {detect_params.dist_crit} degrees"
),
},
),
},
{
"std_name": "atmosphere_upward_relative_vorticity",
"long_name": "Atmosphere Upward Relative Vorticity",
"tstorms_name": "vort850",
"filename": self._in_file("vort"),
"cellmethod": cf.CellMethod(
"area",
"maximum",
qualifiers={
"comment": (
f"lesser circle of radius {detect_params.dist_crit} degrees"
),
},
),
},
]
# Initialise variable metadata as empty dict to populate
self._variable_metadata = {}
for output in var_outputs:
# Check input file exists to read from
if not os.path.exists(output["filename"]):
errmsg = f"Input file '{output['filename']}' not found."
raise FileNotFoundError(errmsg)
var_name = output["std_name"]
tstorms_name = output["tstorms_name"]
# Get the variable field from the netcdf file
fields = cf.read(output["filename"], select=f"ncvar%{tstorms_name}") # type: ignore[operator]
if not fields:
msg = (
f"Variable '{tstorms_name}' not found "
f"in input file {output['filename']}."
)
raise ValueError(msg)
field = fields[0]
# Read and store the relevant metadata
# Override wind speed as reading from u file but this is zonal velocity
self._variable_metadata[var_name] = TCTrackerMetadata(
{
"standard_name": (
output["std_name"]
if var_name == "wind_speed"
else field.get_property("standard_name", output["std_name"])
),
"long_name": (
output["long_name"]
if var_name == "wind_speed"
else field.get_property("long_name", output["long_name"])
),
"units": field.get_property("units", "unknown"),
}
)
self._variable_metadata[var_name].constructs = [output["cellmethod"]]
def _set_time_metadata(self):
"""
Extract time metadata from u_ref input and set attribute.
Calendar type, units, start time, and end time extracted from u_ref as for
TSTORMS by identifying the unlimited dimension and fetching coordinate
attributes. Assumes other files match.
If file not found or variable cannot be identified errors are raised.
The calendar defaults to Julian as for TSTORMS.
Raises
------
KeyError
If no unlimited dimension or corresponding coordinate is found in the u
input file.
ValueError
If multiple unlimited dimensions are found in the u input file.
ValueError
If the 'units' attribute is missing for the time coordinate.
Warnings
--------
UserWarning
If the 'calendar' attribute is missing for the coordinate variable.
Defaults to 'julian'.
"""
input_u_file = self._in_file("u")
with Dataset(input_u_file, "r") as nc_file:
unlimited_dims = [
dim
for dim in nc_file.dimensions
if nc_file.dimensions[dim].isunlimited()
]
if not unlimited_dims:
errmsg = (
"No unlimited dimension found in the u_ref NetCDF file "
"to set calendar."
)
raise KeyError(errmsg)
if len(unlimited_dims) > 1:
errmsg = (
"TSTORMS expects only a single unlimited variable in NetCDF "
"files corresponding to the time dimension. "
f"Multiple found: {unlimited_dims}."
)
raise ValueError(errmsg)
unlimited_dim = unlimited_dims[0]
if unlimited_dim not in nc_file.variables:
errmsg = (
"Coordinate variable for unlimited dimension "
f"'{unlimited_dim}' not found."
)
raise KeyError(errmsg)
coord_var = nc_file.variables[unlimited_dim]
units = getattr(coord_var, "units", None)
if units is None:
msg = "The 'units' attribute is required for the time coordinate."
raise ValueError(msg)
calendar = getattr(coord_var, "calendar", None)
if calendar is None:
msg = (
"The 'calendar' attribute is missing for the coordinate variable.\n"
"defaulting to 'julian'"
)
warnings.warn(msg, category=UserWarning, stacklevel=2)
calendar = "julian"
start_time_num = coord_var[0]
end_time_num = coord_var[-1]
self._time_metadata = {
"calendar": calendar,
"units": units,
"start_time": num2date(start_time_num, units=units, calendar=calendar),
"end_time": num2date(end_time_num, units=units, calendar=calendar),
}
[docs]
def read_trajectories(self):
"""
Parse outputs from TSTORMS to list of :class:`tctrack.core.Trajectory`.
This will read from the ``trav`` file output from trajectory stitching, or
``trav_filt`` if filtering was applied (see ``do_filter`` in the
:attr:`tstorms_parameters` attribute).
Returns
-------
trajectories : list[Trajectory]
A list of :class:`tctrack.core.Trajectory` objects.
"""
trajectories = []
current_trajectory_id = 0 # Initialize trajectory ID
# We need time metadata, so set the metadata if not done already
if not self._time_metadata:
self._set_time_metadata()
# Use the trav file to get full data (including vorticity)
# If filtering was applied use filtered file
output_dir = self.tstorms_parameters.output_dir
trav_file = "trav_filt" if self.stitch_parameters.do_filter else "trav"
output_trav_filepath = os.path.join(output_dir, trav_file)
with open(output_trav_filepath, "r") as file:
for line in file:
items = line.split()
if items[0] == "start":
# Start of new trajectory.
# Extract metadata and add Trajectory to dict
current_trajectory_id += 1
time = list(map(int, items[2:6]))
trajectories.append(
Trajectory(
current_trajectory_id,
time,
calendar=self.time_metadata["calendar"],
)
)
# Continue processing ongoing trajectory
# Trajectories indexed from 1 so subtract when assigning
else:
trajectories[current_trajectory_id - 1].add_point(
*self._parse_tstorms_trav_line_to_point(items)
)
return trajectories
@staticmethod
def _parse_tstorms_trav_line_to_point(
line: list[str], variable_names: list[str] | None = None
) -> tuple[list[int], dict[str, int | float]]:
"""
Parse line from TSTORMS trav output into a trajectory data point.
Data point format is that expected by a :class:`tctrack.core.Trajectory`.
Data from a traj file can be parsed by providing a different set of
``variable_names``
Parameters
----------
line : list[str]
A list of strings representing the line split into parts.
variable_names : list[str] | None
List of variable names for the data columns.
Defaults to None and then set to those for a trav file.
Returns
-------
tuple
A tuple containing the time as an integer list of [year, day, month, hour]
and a dict of variables.
"""
if not variable_names:
variable_names = [
"lon",
"lat",
"wind_speed",
"air_pressure_at_mean_sea_level",
"atmosphere_upward_relative_vorticity",
]
return_vars: dict[str, int | float] = {}
return_vars.update(
{
name: float(value)
for name, value in zip(variable_names, line[:-4], strict=True)
}
)
time = list(map(int, line[-4:]))
return time, return_vars
[docs]
def run_tracker(self, input_files: str | Iterable[str], output_file: str):
"""Run TSTORMS tracker to obtain tropical cyclone track trajectories.
This first runs :meth:`detect` to get TC candidates at each time. Then
these are combined into trajectories using :meth:`stitch`.
The output is then saved as a CF-compliant NetCDF trajectory file.
Arguments
---------
input_files : str | Iterable[str]
A (list of) file path(s) containing NetCDF input data to use in the tracker.
output_file : str
Filename to which the tropical cyclone trajectories are saved.
Raises
------
FileNotFoundError
- If the TSTORMS executables cannot be found.
RuntimeError
If the TSTORMS commands return a non-zero exit code.
Examples
--------
To set the parameters, instantiate a :class:`TSTORMSTracker` instance and run
`run_tracker()`:
>>> input_files = [...]
>>> tstorms_params = TSTORMSBaseParameters(...)
>>> detect_params = TSTORMSDetectParameters(...)
>>> stitch_params = TSTORMSStitchParameters(...)
>>> my_tracker = TSTORMSTracker(tstorms_params, detect_params, stitch_params)
>>> my_tracker.run_tracker(input_files, "trajectories.nc")
"""
self.set_input_files(input_files)
self.detect()
self.stitch()
self.to_netcdf(output_file)