API#

earth2mip.diagnostics module#

class earth2mip.diagnostics.Diagnostics(group, domain, grid, diagnostic, lat, lon, device)#

Bases: object

get_dimensions()#
get_dtype()#
get_variables()#
update()#
class earth2mip.diagnostics.Raw(group, domain, grid, diagnostic, lat, lon, device)#

Bases: Diagnostics

get_dimensions()#
get_dtype()#
update(output, time_index, batch_id, batch_size)#

earth2mip.ensemble_utils module#

class earth2mip.ensemble_utils.GaussianRandomFieldS2(nlat, alpha=2.0, tau=3.0, sigma=None, radius=1.0, grid='equiangular', dtype=torch.float32)#

Bases: Module

cuda(*args, **kwargs)#

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters:

device (int, optional) – if specified, all parameters will be copied to that device

Returns:

self

Return type:

Module

forward(N, xi=None)#

Sample random functions from a spherical GRF.

Parameters:

  • N (int) – Number of functions to sample.

  • xi (torch.Tensor, default is None) – Noise is a complex tensor of size (N, nlat, nlat+1). If None, new Gaussian noise is sampled. If xi is provided, N is ignored.

  • Output –

  • ------- –

  • u (torch.Tensor) – N random samples from the GRF returned as a tensor of size (N, nlat, 2*nlat) on a equiangular grid.

to(*args, **kwargs)#

Moves and/or casts the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)
to(dtype, non_blocking=False)
to(tensor, non_blocking=False)
to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

  • device (torch.device) – the desired device of the parameters and buffers in this module

  • dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

  • tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

  • memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
earth2mip.ensemble_utils.brown_noise(shape, reddening=2)#
earth2mip.ensemble_utils.generate_bred_vector(x, model, noise_amplitude, time=None, integration_steps=40, inflate=False)#
Return type:

Tensor

earth2mip.ensemble_utils.generate_noise_correlated(shape, *, reddening, device, noise_amplitude)#
earth2mip.ensemble_utils.generate_noise_grf(shape, grid, alpha, sigma, tau)#

earth2mip.filesystem module#

earth2mip.filesystem.download_cached(path, recursive=False)#
Return type:

str

earth2mip.filesystem.glob(pattern, maxdepth=1)#
Return type:

List[str]

earth2mip.filesystem.ls(path)#
earth2mip.filesystem.open(path, mode='r')#
earth2mip.filesystem.pipe(dest, value)#

Save string to dest

earth2mip.forecast_metrics_io module#

Routines for reading and writing forecast metrics to a directory of csv files.

The csv files contain the records:

initial_time_iso, lead_time_hours, channel, metric, value
2022-01-01T00:00:00,24,t2m,rmse,25.6
earth2mip.forecast_metrics_io.read_metrics(directory)#

Reads all csv files in the given directory and returns a pandas Series containing all the metric values.

Return type:

Series

earth2mip.forecast_metrics_io.write_metric(f, initial_time, lead_time, channel, metric, value)#

Writes a single metric value to the given file object in csv format.

Return type:

None

earth2mip.forecasts module#

Forecast abstractions

A forecast is a discrete array of (n_initial_times, n_lead_times). However because a forecast evolves forward in time, and we do not store the whole forecast necessarily, algorithms in fcn-mip should access n_lead_times in sequential order. This is the purpose of the abstractions here.

class earth2mip.forecasts.Forecast(*args, **kwargs)#

Bases: Protocol

property channel_names: List[str]#
property grid: LatLonGrid#
class earth2mip.forecasts.Persistence(observations)#

Bases: object

persistence forecast. This forecast always returns the initial condition.

Yields (channel, lat, lon)

property channel_names#
class earth2mip.forecasts.TimeLoopForecast(time_loop, times, data_source)#

Bases: Forecast

Wrap an fcn-mip TimeLoop object as a forecast

property channel_names#
property grid#
class earth2mip.forecasts.XarrayForecast(ds, fields, times, device)#

Bases: Forecast

Turn an xarray into a forecast-like dataset

property channel_names#
property grid#

earth2mip.geo_operator module#

class earth2mip.geo_operator.GeoOperator(*args, **kwargs)#

Bases: Protocol

Geo Operator

This is the most primative functional of Earth-2 MIP which represents a operators on geospatial fields. This implies the following two requirements:

  1. The operation must define in and out channel variables representing the

    fields in the input/output arrays.

  2. The operation must define the in and out grid schemas.

Many auto-gressive models can be represented as a GeoOperator and can maintain a internal state. Diagnostic models must be a GeoOperator by definition.

Warning

Geo Function is a concept not full adopted in Earth-2 MIP and is being adopted progressively.

property in_channel_names: list[str]#
property in_grid: LatLonGrid#
property out_channel_names: list[str]#
property out_grid: LatLonGrid#

earth2mip.geometry module#

Routines for working with geometry

earth2mip.geometry.bilinear(data, dims, source_coords, target_coords)#
earth2mip.geometry.get_batch_size(data)#
earth2mip.geometry.get_bounds_window(geom, lat, lon)#
earth2mip.geometry.select_space(data, lat, lon, domain)#

earth2mip.grid module#

class earth2mip.grid.LatLonGrid(lat, lon)#

Bases: object

lat: List[float]#
lon: List[float]#
property shape#
earth2mip.grid.equiangular_lat_lon_grid(nlat, nlon, includes_south_pole=True)#

A regular lat-lon grid

Lat is ordered from 90 to -90. Includes -90 and only if if includes_south_pole is True. Lon is ordered from 0 to 360. includes 0, but not 360.

Return type:

LatLonGrid

earth2mip.grid.from_enum(grid_enum)#
Return type:

LatLonGrid

earth2mip.inference_ensemble module#

earth2mip.inference_medium_range module#

earth2mip.loaders module#

class earth2mip.loaders.LoaderProtocol(*args, **kwargs)#

Bases: Protocol

earth2mip.loaders.torchscript(package, pretrained=True)#

load a checkpoint into a model

earth2mip.make_job module#

earth2mip.make_job.get_time(times)#
earth2mip.make_job.get_time_s2s_calibration()#
earth2mip.make_job.get_times_2018()#
earth2mip.make_job.get_times_s2s_test()#
earth2mip.make_job.main(model, config, output)#

earth2mip.model_registry module#

Create-read-update-delete (CRUD) operations for the FCN model registry

The location of the registry is configured using config.MODEL_REGISTRY. Both s3:// and local paths are supported.

The top-level structure of the registry is like this:

afno_26ch_v/
baseline_afno_26/
gfno_26ch_sc3_layers8_tt64/
hafno_baseline_26ch_edim512_mlp2/
modulus_afno_20/
sfno_73ch/
tfno_no-patching_lr5e-4_full_epochs/

The name of the model is the folder name. Each of these folders has the following structure:

sfno_73ch/about.txt            # optional information (e.g. source path)
sfno_73ch/global_means.npy
sfno_73ch/global_stds.npy
sfno_73ch/weights.tar          # model checkpoint
sfno_73ch/metadata.json

The metadata.json file contains data necessary to use the model for forecasts:

{
    "architecture": "sfno_73ch",
    "n_history": 0,
    "grid": "721x1440",
    "in_channels": [
        0,
        1
    ],
    "out_channels": [
        0,
        1
    ]
}

Its schema is provided by the earth2mip.schema.Model.

The checkpoint file weights.tar should have a dictionary of model weights and parameters in the model_state key. For backwards compatibility with FCN checkpoints produced as of March 1, 2023 the keys should include prefixed module. prefix. This checkpoint format may change in the future.

Scoring FCNs under active development#

One can use fcn-mip to score models not packaged in fcn-mip using a metadata file like this:

{
    "architecture": "pickle",
    ...
}

This will load weights.tar using torch.load. This is not recommended for long-time archival of model checkpoints but does allow scoring models under active development. Once a reasonable skill is achieved the model’s source code can be stabilized and packaged within fcn-mip for long-term archival.

earth2mip.model_registry.DLWPPackage(root, seperator)#
earth2mip.model_registry.FCNPackage(root, seperator)#
earth2mip.model_registry.FCNv2Package(root, seperator)#
class earth2mip.model_registry.ModelRegistry(path)#

Bases: object

SEPERATOR: str = '/'#
get_builtin_model(name)#

Built in models that have globally buildable packages

get_center_path(name)#
get_metadata(name)#
Return type:

Model

get_model(name)#
get_model_path(name)#
get_path(name, *args)#
get_scale_path(name)#
get_weight_path(name)#
list_models()#
put_metadata(name, metadata)#
earth2mip.model_registry.NGCDiagnosticPackage(root, seperator, name)#
class earth2mip.model_registry.Package(root, seperator)#

Bases: object

A model package

Simple file system operations and quick metadata access

get(path, recursive=False)#
metadata()#
Return type:

Model

earth2mip.model_registry.PanguPackage(root, seperator)#
earth2mip.model_registry.download_ngc_package(root, url, zip_file)#

earth2mip.netcdf module#

Routines to save domains to a netCDF file

earth2mip.netcdf.initialize_netcdf(nc, domains, grid, n_ensemble, device)#
Return type:

List[List[Diagnostics]]

earth2mip.netcdf.update_netcdf(data, total_diagnostics, domains, batch_id, time_count, grid, channel_names_of_data)#

earth2mip.regrid module#

class earth2mip.regrid.Identity(*args, **kwargs)#

Bases: Module

forward(x)#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class earth2mip.regrid.RegridLatLon(src_grid, dest_grid)#

Bases: Module

forward(x)#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class earth2mip.regrid.TempestRegridder(file_path)#

Bases: Module

forward(x)#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

to(device)#

Moves and/or casts the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)
to(dtype, non_blocking=False)
to(tensor, non_blocking=False)
to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

  • device (torch.device) – the desired device of the parameters and buffers in this module

  • dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

  • tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

  • memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
earth2mip.regrid.get_regridder(src, dest)#
Return type:

Module

earth2mip.schema module#

class earth2mip.schema.InitialConditionSource(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)#

Bases: Enum

cds: str = 'cds'#
era5: str = 'era5'#
gfs: str = 'gfs'#
hrmip: str = 'hrmip'#
ifs: str = 'ifs'#
class earth2mip.schema.WeatherEvent(**data)#

Bases: BaseModel

domains: List[Union[Window, CWBDomain, MultiPoint]]#
properties: WeatherEventProperties#

earth2mip.score_ensemble_outputs module#

earth2mip.score_ensemble_outputs.main(input_path, output_path=None, time_averaging_window='', score=True, save_ensemble=False)#
Return type:

None

earth2mip.score_ensemble_outputs.open_ensemble(path, group)#
earth2mip.score_ensemble_outputs.open_verification(time)#
earth2mip.score_ensemble_outputs.read_weather_event(dir)#
earth2mip.score_ensemble_outputs.save_dataset(out, path)#

earth2mip.time module#

earth2mip.time.convert_to_datetime(time)#
Return type:

datetime

earth2mip.time.datetime_to_timestamp(time)#
Return type:

float

earth2mip.time_collection module#

earth2mip.time_loop module#

class earth2mip.time_loop.TimeLoop(*args, **kwargs)#

Bases: Protocol

Abstract protocol that a custom time loop must follow

This is a callable which yields time and output information. Some attributes are required to define the input and output data required.

The expectation is that this class and the data passed to it are on the same device. While torch modules can be moved between devices easily, this is not true for all frameworks.

in_channel_names#
out_channel_names#
grid#
n_history_levels#
history_time_step#
time_step#
device#
device: device#
dtype: dtype = torch.float32#
grid: LatLonGrid#
history_time_step: timedelta = datetime.timedelta(0)#
in_channel_names: List[str]#
n_history_levels: int = 1#
out_channel_names: List[str]#
time_step: timedelta#

earth2mip.weather_events module#

class earth2mip.weather_events.CWBDomain(**data)#

Bases: BaseModel

diagnostics: List[Diagnostic]#
name: str#
path: str#
type: Literal['CWBDomain']#
class earth2mip.weather_events.Diagnostic(**data)#

Bases: BaseModel

channels: List[str]#
function: str#
nbins: int#
type: str#
class earth2mip.weather_events.InitialConditionSource(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)#

Bases: Enum

cds: str = 'cds'#
era5: str = 'era5'#
gfs: str = 'gfs'#
hrmip: str = 'hrmip'#
ifs: str = 'ifs'#
class earth2mip.weather_events.MultiPoint(**data)#

Bases: BaseModel

diagnostics: List[Diagnostic]#
lat: List[float]#
lon: List[float]#
name: str#
type: Literal['MultiPoint']#
class earth2mip.weather_events.WeatherEvent(**data)#

Bases: BaseModel

domains: List[Union[Window, CWBDomain, MultiPoint]]#
properties: WeatherEventProperties#
class earth2mip.weather_events.WeatherEventProperties(**data)#

Bases: BaseModel

netcdf#

load the initial conditions from this path if given

initial_condition_source: InitialConditionSource#
name: str#
netcdf: str#
restart: str#
start_time: Optional[datetime]#
class earth2mip.weather_events.Window(**data)#

Bases: BaseModel

diagnostics: List[Diagnostic]#
lat_max: float#
lat_min: float#
lon_max: float#
lon_min: float#
name: str#
type: Literal['Window']#
earth2mip.weather_events.list_()#
earth2mip.weather_events.read(forecast_name)#
Return type:

WeatherEvent