earth2studio.models.dx.ClimateNet

class earth2studio.models.dx.ClimateNet(core_model, center, scale)

Climate Net diagnostic model, built into Earth2Studio. This model can be used to create prediction labels for tropical cyclones and atmospheric rivers from a set of three atmospheric variables on a quater degree resolution equirectangular grid. It produces three non-standard output channels climnet_bg, climnet_tc and climnet_ar representing background label, tropical cyclone and atmospheric river labels.

Note

This model and checkpoint are from Prabhat et al. 2021. For more information see the following references:

• https://doi.org/10.5194/gmd-14-107-2021

• andregraubner/ClimateNet

Parameters:

• core_model (torch.nn.Module) – Core pytorch model

• center (torch.Tensor) – Model center normalization tensor of size [20,1,1]

• scale (torch.Tensor) – Model scale normalization tensor of size [20,1,1]

__call__(x, coords)

Forward pass of diagnostic

Parameters:

• x (Tensor)

• coords (OrderedDict[str, ndarray])

Return type:

tuple[Tensor, OrderedDict[str, ndarray]]

classmethod load_default_package()

Default pre-trained climatenet model package from Nvidia model registry

Return type:

Package

classmethod load_model(package)

Load diagnostic from package

Parameters:

package (Package)

Return type:

DiagnosticModel