earth2mip.networks package#

earth2mip.networks.dlwp module#

class earth2mip.networks.dlwp.DLWPInference(dlwp, center, scale)#

Bases: Module

Parameters:

  • center (array) –

  • scale (array) –

property channel_names#
cuda(device=None)#

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

property device: device#
property dtype: dtype#
property grid: LatLonGrid#
history_time_step = datetime.timedelta(seconds=21600)#
property in_channel_names#
property n_history#
n_history_levels = 2#
normalize(x)#
property out_channel_names#
time_step = datetime.timedelta(seconds=21600)#
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)
unnormalize(x)#
earth2mip.networks.dlwp.load(package, *, pretrained=True, device='cuda')#
Parameters:

package (Package) –

earth2mip.networks.fcn module#

FCN adapter from Modulus

earth2mip.networks.fcn.load(package, *, pretrained=True, device='cuda')#

earth2mip.networks.fcnv2_sm module#

FCN v2 Small adapter

This model is an outdated version of FCN v2 (SFNO), a more recent one is present in Modulus.

earth2mip.networks.fcnv2_sm.load(package, *, pretrained=True, device='cuda')#

earth2mip.networks.graphcast module#

earth2mip.networks.graphcast.load_time_loop(package, pretrained=True, device='cuda:0')#
Return type:

TimeStepperLoop

earth2mip.networks.graphcast.load_time_loop_operational(package, pretrained=True, device='cuda:0')#
Return type:

TimeStepperLoop

earth2mip.networks.graphcast.load_time_loop_small(package, pretrained=True, device='cuda:0')#
Return type:

TimeStepperLoop

earth2mip.networks.pangu module#

Pangu Weather adapter

adapted from https://raw.githubusercontent.com/ecmwf-lab/ai-models-panguweather/main/ai_models_panguweather/model.py

# (C) Copyright 2023 European Centre for Medium-Range Weather Forecasts. # This software is licensed under the terms of the Apache Licence Version 2.0 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. # In applying this licence, ECMWF does not waive the privileges and immunities # granted to it by virtue of its status as an intergovernmental organisation # nor does it submit to any jurisdiction.

class earth2mip.networks.pangu.PanguInference(model_6, model_24)#

Bases: Module

Parameters:
cuda(device=None)#

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

property device: device#
dtype: dtype = torch.float32#
property grid: LatLonGrid#
history_time_step = datetime.timedelta(0)#
property in_channel_names#
property n_history#
n_history_levels = 1#
normalize(x)#
property out_channel_names#
run_steps_with_restart(x, n, normalize=True, time=None)#

Yield (time, unnormalized data, restart) tuples

restart = (time, unnormalized data)

time_step = datetime.timedelta(seconds=21600)#
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)
class earth2mip.networks.pangu.PanguStacked(model)#

Bases: object

Parameters:

model (PanguWeather) –

channel_names()#
forward(x)#
to()#
class earth2mip.networks.pangu.PanguWeather(path)#

Bases: object

area = [90, 0, -90, 360]#
download_files = ['pangu_weather_24.onnx', 'pangu_weather_6.onnx']#
download_url = 'https://get.ecmwf.int/repository/test-data/ai-models/pangu-weather/{file}'#
expver = 'pguw'#
grid = [0.25, 0.25]#
param_level_pl = (['z', 'q', 't', 'u', 'v'], [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50])#
param_sfc = ['msl', 'u10m', 'v10m', 't2m']#
earth2mip.networks.pangu.load(package, *, pretrained=True, device="doesn't matter")#

Load the sub-stepped pangu weather inference

earth2mip.networks.pangu.load_24(package, *, pretrained=True, device='cuda:0')#

Load a 24 hour time-step pangu weather

earth2mip.networks.pangu.load_6(package, *, pretrained=True, device='cuda:0')#

Load a 6 hour time-step pangu weather

earth2mip.networks.pangu.load_single_model(package, *, time_step_hours=24, pretrained=True, device='cuda:0')#

Load a single time-step pangu weather

Parameters:

time_step_hours (int) –

Module contents#

earth2mip.networks.get_model(model, registry=<earth2mip.model_registry.ModelRegistry object>, device='cpu', metadata=None)#

Function to construct an inference model and load the appropriate checkpoints from the model registry

Parameters:

  • model (The model name to open in the registry. If a url is passed (e.g.) – s3://bucket/model), then this location will be opened directly. Supported urls protocols include s3:// for PBSS access, and file:// for local files.

  • registry (A model registry object. Defaults to the global model registry) –

  • metadata (If provided, this model metadata will be used to load the model.) – By default this will be loaded from the file metadata.json in the model package.

  • device (the device to load on, by default the 'cpu') –

Return type:

Inference model