API Reference

This page is automatically generated via mkdocstrings from the current codebase.

Training & Inference Entry Points

train

train

apply_config_overrides(overrides, configs_by_root)

Apply strict nested config overrides in-place.

All path segments must already exist; this prevents silent typos.

build_dataset(data_config_path, ds_cfg)

Build and return dataset.

Parameters:

Name	Type	Description	Default
data_config_path	str	Path to an input or output file.	required
ds_cfg	dict[str, Any]	Configuration dictionary or section.	required

Returns:

Type	Description
SurfaceTempPatchLightDataset | SurfaceTempPatch4BandsLightDataset	SurfaceTempPatchLightDataset | SurfaceTempPatch4BandsLightDataset: Computed output value.

build_wandb_logger(training_cfg, model)

Build and return wandb logger.

Parameters:

Name	Type	Description	Default
training_cfg	dict[str, Any]	Configuration dictionary or section.	required
model	LightningModule	Input value.	required

Returns:

Name	Type	Description
WandbLogger	WandbLogger	Computed output value.

dump_yaml(path, payload)

Write yaml data to disk preserving key order.

load_yaml(path)

Load and return yaml data.

Parameters:

Name	Type	Description	Default
path	str	Path to an input or output file.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Dictionary containing computed outputs.

main(model_config_path='configs/model_config.yaml', data_config_path='configs/data_config.yaml', training_config_path='configs/training_config.yaml', overrides=None)

Run the script entry point.

Parameters:

Name	Type	Description	Default
model_config_path	str	Path to an input or output file.	'configs/model_config.yaml'
data_config_path	str	Path to an input or output file.	'configs/data_config.yaml'
training_config_path	str	Path to an input or output file.	'configs/training_config.yaml'
overrides	list[str] | None	Optional config overrides from CLI.	None

Returns:

Name	Type	Description
None	None	No value is returned.

parse_args()

Parse command-line arguments for this script.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Type	Description
Namespace	argparse.Namespace: Computed output value.

parse_config_override(raw_override)

Parse one CLI config override expression.

Expected format: .= where config_root is one of: data, training, model.

resolve_dataset_variant(ds_cfg, data_config_path)

Resolve and validate dataset variant.

Parameters:

Name	Type	Description	Default
ds_cfg	dict[str, Any]	Configuration dictionary or section.	required
data_config_path	str	Path to an input or output file.	required

Returns:

Name	Type	Description
str	str	Computed scalar output.

resolve_global_rank()

Resolve and validate global rank.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
int	int	Computed scalar output.

resolve_model_type(model_cfg)

Resolve and validate model type.

Parameters:

Name	Type	Description	Default
model_cfg	dict[str, Any]	Configuration dictionary or section.	required

Returns:

Name	Type	Description
str	str	Computed scalar output.

resolve_resume_ckpt_path(model_cfg)

Resolve and validate resume ckpt path.

Parameters:

Name	Type	Description	Default
model_cfg	dict[str, Any]	Configuration dictionary or section.	required

Returns:

Type	Description
str | None	str | None: Computed output value.

resolve_wandb_watch_mode(wandb_cfg)

Resolve and validate wandb watch mode.

Parameters:

Name	Type	Description	Default
wandb_cfg	dict[str, Any]	Configuration dictionary or section.	required

Returns:

Type	Description
str | None	str | None: Computed output value.

upload_configs_to_wandb(logger, config_paths)

Upload configs to wandb to experiment tracking.

Parameters:

Name	Type	Description	Default
logger	WandbLogger	Logger instance used for experiment tracking.	required
config_paths	list[str]	Path to an input or output file.	required

Returns:

Name	Type	Description
None	None	No value is returned.

inference

inference

build_datamodule(dataset, data_cfg, training_cfg)

Build and return datamodule.

Parameters:

Name	Type	Description	Default
dataset	Dataset	Input value.	required
data_cfg	dict[str, Any]	Configuration dictionary or section.	required
training_cfg	dict[str, Any]	Configuration dictionary or section.	required

Returns:

Name	Type	Description
DepthTileDataModule	DepthTileDataModule	Computed output value.

build_dataset(data_config_path, ds_cfg)

Build and return dataset.

Parameters:

Name	Type	Description	Default
data_config_path	str	Path to an input or output file.	required
ds_cfg	dict[str, Any]	Configuration dictionary or section.	required

Returns:

Type	Description
SurfaceTempPatchLightDataset | SurfaceTempPatch4BandsLightDataset	SurfaceTempPatchLightDataset | SurfaceTempPatch4BandsLightDataset: Computed output value.

build_model(model_config_path, data_config_path, training_config_path, model_cfg, datamodule)

Build and return model.

Parameters:

Name	Type	Description	Default
model_config_path	str	Path to an input or output file.	required
data_config_path	str	Path to an input or output file.	required
training_config_path	str	Path to an input or output file.	required
model_cfg	dict[str, Any]	Configuration dictionary or section.	required
datamodule	DepthTileDataModule	Input value.	required

Returns:

Name	Type	Description
PixelDiffusionConditional	PixelDiffusionConditional	Computed output value.

build_random_batch(model, data_cfg, batch_size, height, width, device)

Build and return random batch.

Parameters:

Name	Type	Description	Default
model	PixelDiffusionConditional	Input value.	required
data_cfg	dict[str, Any]	Configuration dictionary or section.	required
batch_size	int	Size/count parameter.	required
height	int	Size/count parameter.	required
width	int	Size/count parameter.	required
device	device	Target device for tensor operations.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Dictionary containing computed outputs.

choose_device(device_arg)

Choose and return device.

Parameters:

Name	Type	Description	Default
device_arg	str	Input value.	required

Returns:

Type	Description
device	torch.device: Computed output value.

load_yaml(path)

Load and return yaml data.

Parameters:

Name	Type	Description	Default
path	str	Path to an input or output file.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Dictionary containing computed outputs.

main()

Run the script entry point.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
None	None	No value is returned.

pretty_shape(value)

Return a compact human-readable shape/type description.

Parameters:

Name	Type	Description	Default
value	Any	Input value.	required

Returns:

Name	Type	Description
str	str	Computed scalar output.

resolve_checkpoint_path(ckpt_override, model_cfg)

Resolve and validate checkpoint path.

Parameters:

Name	Type	Description	Default
ckpt_override	str | None	Input value.	required
model_cfg	dict[str, Any]	Configuration dictionary or section.	required

Returns:

Type	Description
str | None	str | None: Computed output value.

resolve_dataset_variant(ds_cfg, data_config_path)

Resolve and validate dataset variant.

Parameters:

Name	Type	Description	Default
ds_cfg	dict[str, Any]	Configuration dictionary or section.	required
data_config_path	str	Path to an input or output file.	required

Returns:

Name	Type	Description
str	str	Computed scalar output.

resolve_model_type(model_cfg)

Resolve and validate model type.

Parameters:

Name	Type	Description	Default
model_cfg	dict[str, Any]	Configuration dictionary or section.	required

Returns:

Name	Type	Description
str	str	Computed scalar output.

run_predict_once(model, batch, include_intermediates)

Compute run predict once and return the result.

Parameters:

Name	Type	Description	Default
model	PixelDiffusionConditional	Input value.	required
batch	dict[str, Any]	Input value.	required
include_intermediates	bool	Boolean flag controlling behavior.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Dictionary containing computed outputs.

to_device(batch, device)

Move tensor values in a batch dictionary to the target device.

Parameters:

Name	Type	Description	Default
batch	dict[str, Any]	Input value.	required
device	device	Target device for tensor operations.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Dictionary containing computed outputs.

Data

data.datamodule

data.datamodule

DepthTileDataModule

Bases: LightningDataModule

Lightning DataModule that builds train and validation dataloaders.

__init__(*, dataset, val_dataset=None, dataloader_cfg=None, val_fraction=0.2, seed=7)

Initialize DepthTileDataModule with configured parameters.

Parameters:

Name	Type	Description	Default
dataset	Dataset	Input value.	required
val_dataset	Dataset | None	Input value.	None
dataloader_cfg	dict[str, Any] | None	Configuration dictionary or section.	None
val_fraction	float	Input value.	0.2
seed	int	Input value.	7

Returns:

Name	Type	Description
None	None	No value is returned.

setup(stage=None)

Compute setup and return the result.

Parameters:

Name	Type	Description	Default
stage	str | None	Input value.	None

Returns:

Name	Type	Description
None	None	No value is returned.

train_dataloader()

Return the training dataloader from the attached datamodule.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
DataLoader	DataLoader	Computed output value.

val_dataloader()

Return the validation dataloader from the attached datamodule.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
DataLoader	DataLoader	Computed output value.

data.dataset_temp_v1

data.dataset_temp_v1

SurfaceTempPatchLightDataset

Bases: Dataset

Dataset that loads single-band surface temperature patches.

__getitem__(idx)

Load and return one sample for the given index.

Parameters:

Name	Type	Description	Default
idx	int	Zero-based index for selecting a sample or batch.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Dictionary containing computed outputs.

__init__(*, csv_path, split='all', mask_fraction=0.0, mask_patch_min=3, mask_patch_max=9, enable_transform=False, x_return_mode='currupted_plus_mask', return_info=False, return_coords=False, nan_fill_value=0.0, valid_from_fill_value=True, split_seed=7, val_fraction=0.2)

Initialize SurfaceTempPatchLightDataset with configured parameters.

Parameters:

Name	Type	Description	Default
csv_path	str | Path	Path to an input or output file.	required
split	str	Input value.	'all'
mask_fraction	float	Mask tensor controlling valid or known pixels.	0.0
mask_patch_min	int	Mask tensor controlling valid or known pixels.	3
mask_patch_max	int	Mask tensor controlling valid or known pixels.	9
enable_transform	bool	Boolean flag controlling behavior.	False
x_return_mode	str	Input value.	'currupted_plus_mask'
return_info	bool	Boolean flag controlling behavior.	False
return_coords	bool	Boolean flag controlling behavior.	False
nan_fill_value	float	Input value.	0.0
valid_from_fill_value	bool	Boolean flag controlling behavior.	True
split_seed	int	Input value.	7
val_fraction	float	Input value.	0.2

Returns:

Name	Type	Description
None	None	No value is returned.

__len__()

Return the number of available samples.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
int	int	Computed scalar output.

from_config(config_path='configs/data_config.yaml', *, split='all') classmethod

Compute from config and return the result.

Parameters:

Name	Type	Description	Default
config_path	str	Path to an input or output file.	'configs/data_config.yaml'
split	str	Input value.	'all'

Returns:

Type	Description
'SurfaceTempPatchLightDataset'	'SurfaceTempPatchLightDataset': Computed output value.

data.dataset_4bands

data.dataset_4bands

SurfaceTempPatch4BandsLightDataset

Bases: Dataset

Dataset that loads EO-conditioned multi-band depth patches.

__getitem__(idx)

Load and return one sample for the given index.

Parameters:

Name	Type	Description	Default
idx	int	Zero-based index for selecting a sample or batch.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Dictionary containing computed outputs.

__init__(*, csv_path, split='all', mask_fraction=0.0, mask_patch_min=3, mask_patch_max=9, enable_transform=False, x_return_mode='currupted_plus_mask', return_info=False, return_coords=False, nan_fill_value=0.0, valid_from_fill_value=True, eo_random_scale_enabled=False, eo_speckle_noise_enabled=False, split_seed=7, val_fraction=0.2)

Initialize SurfaceTempPatch4BandsLightDataset with configured parameters.

Parameters:

Name	Type	Description	Default
csv_path	str | Path	Path to an input or output file.	required
split	str	Input value.	'all'
mask_fraction	float	Mask tensor controlling valid or known pixels.	0.0
mask_patch_min	int	Mask tensor controlling valid or known pixels.	3
mask_patch_max	int	Mask tensor controlling valid or known pixels.	9
enable_transform	bool	Boolean flag controlling behavior.	False
x_return_mode	str	Input value.	'currupted_plus_mask'
return_info	bool	Boolean flag controlling behavior.	False
return_coords	bool	Boolean flag controlling behavior.	False
nan_fill_value	float	Input value.	0.0
valid_from_fill_value	bool	Boolean flag controlling behavior.	True
eo_random_scale_enabled	bool	Boolean flag controlling behavior.	False
eo_speckle_noise_enabled	bool	Boolean flag controlling behavior.	False
split_seed	int	Input value.	7
val_fraction	float	Input value.	0.2

Returns:

Name	Type	Description
None	None	No value is returned.

__len__()

Return the number of available samples.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
int	int	Computed scalar output.

from_config(config_path='configs/data_config.yaml', *, split='all') classmethod

Compute from config and return the result.

Parameters:

Name	Type	Description	Default
config_path	str	Path to an input or output file.	'configs/data_config.yaml'
split	str	Input value.	'all'

Returns:

Type	Description
'SurfaceTempPatch4BandsLightDataset'	'SurfaceTempPatch4BandsLightDataset': Computed output value.

data.dataset_to_disk

data.dataset_to_disk

to_disk(*, root_dir, output_dir, bands, variable='thetao', edge_size=128, enforce_validity=True, max_nodata_fraction=0.25, nan_fill_value=0.0, valid_fraction_threshold=0.25, val_fraction=0.2, split_seed=42, flush_every=500)

Export dataset patches and index metadata to disk.

Parameters:

Name	Type	Description	Default
root_dir	str | Path	Input value.	required
output_dir	str | Path	Input value.	required
bands	Sequence[int]	Input value.	required
variable	str	Input value.	'thetao'
edge_size	int	Input value.	128
enforce_validity	bool	Boolean flag controlling behavior.	True
max_nodata_fraction	float	Input value.	0.25
nan_fill_value	float	Input value.	0.0
valid_fraction_threshold	float	Input value.	0.25
val_fraction	float	Input value.	0.2
split_seed	int	Input value.	42
flush_every	int	Input value.	500

Returns:

Name	Type	Description
Path	Path	Computed output value.

data.assign_window_split

data.assign_window_split

assign_split(input_csv, output_csv, *, val_fraction, seed)

Assign split according to configured rules.

Parameters:

Name	Type	Description	Default
input_csv	Path	Input value.	required
output_csv	Path	Input value.	required
val_fraction	float	Input value.	required
seed	int	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

main()

Run the script entry point.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
None	None	No value is returned.

Model Core

models.difFF.PixelDiffusion

models.difFF.PixelDiffusion

PixelDiffusionConditional

Bases: LightningModule

Lightning module that trains and samples conditional pixel diffusion.

__init__(datamodule=None, generated_channels=1, condition_channels=1, condition_mask_channels=1, condition_include_eo=False, condition_use_valid_mask=True, clamp_known_pixels=True, mask_loss_with_valid_pixels=False, parameterization='epsilon', num_timesteps=1000, noise_schedule='linear', noise_beta_start=0.0001, noise_beta_end=0.02, unet_dim=64, unet_dim_mults=(1, 2, 4, 8), unet_with_time_emb=True, unet_output_mean_scale=False, unet_residual=False, coord_conditioning_enabled=False, coord_encoding='unit_sphere', date_conditioning_enabled=False, date_encoding='day_of_year_sincos', coord_embed_dim=None, batch_size=1, lr=0.001, lr_scheduler_enabled=False, lr_scheduler_monitor='val/loss_ckpt', lr_scheduler_mode='min', lr_scheduler_factor=0.5, lr_scheduler_patience=10, lr_scheduler_threshold=0.0001, lr_scheduler_threshold_mode='rel', lr_scheduler_cooldown=0, lr_scheduler_min_lr=0.0, lr_scheduler_eps=1e-08, lr_warmup_enabled=True, lr_warmup_steps=1000, lr_warmup_start_ratio=0.1, val_inference_sampler='ddpm', val_ddim_num_timesteps=200, val_ddim_eta=0.0, log_intermediates=True, skip_full_reconstruction_in_sanity_check=True, max_full_reconstruction_samples=4, postprocess_gaussian_blur_enabled=False, postprocess_gaussian_blur_sigma=0.35, postprocess_gaussian_blur_kernel_size=3, wandb_verbose=True, log_stats_every_n_steps=1, log_images_every_n_steps=200)

Initialize PixelDiffusionConditional with configured parameters.

Parameters:

Name	Type	Description	Default
datamodule	LightningDataModule | None	Input value.	None
generated_channels	int	Input value.	1
condition_channels	int	Input value.	1
condition_mask_channels	int	Mask tensor controlling valid or known pixels.	1
condition_include_eo	bool	Boolean flag controlling behavior.	False
condition_use_valid_mask	bool	Mask tensor controlling valid or known pixels.	True
clamp_known_pixels	bool	Boolean flag controlling behavior.	True
mask_loss_with_valid_pixels	bool	Mask tensor controlling valid or known pixels.	False
parameterization	str	Input value.	'epsilon'
num_timesteps	int	Step or timestep value.	1000
noise_schedule	str	Input value.	'linear'
noise_beta_start	float	Input value.	0.0001
noise_beta_end	float	Input value.	0.02
unet_dim	int	Input value.	64
unet_dim_mults	tuple[int, ...]	Input value.	(1, 2, 4, 8)
unet_with_time_emb	bool	Boolean flag controlling behavior.	True
unet_output_mean_scale	bool	Boolean flag controlling behavior.	False
unet_residual	bool	Boolean flag controlling behavior.	False
coord_conditioning_enabled	bool	Boolean flag controlling behavior.	False
coord_encoding	str	Input value.	'unit_sphere'
date_conditioning_enabled	bool	Boolean flag controlling behavior.	False
date_encoding	str	Input value.	'day_of_year_sincos'
coord_embed_dim	int | None	Input value.	None
batch_size	int	Size/count parameter.	1
lr	float	Input value.	0.001
lr_scheduler_enabled	bool	Boolean flag controlling behavior.	False
lr_scheduler_monitor	str	Input value.	'val/loss_ckpt'
lr_scheduler_mode	str	Input value.	'min'
lr_scheduler_factor	float	Input value.	0.5
lr_scheduler_patience	int	Input value.	10
lr_scheduler_threshold	float	Input value.	0.0001
lr_scheduler_threshold_mode	str	Input value.	'rel'
lr_scheduler_cooldown	int	Input value.	0
lr_scheduler_min_lr	float	Input value.	0.0
lr_scheduler_eps	float	Input value.	1e-08
lr_warmup_enabled	bool	Boolean flag controlling behavior.	True
lr_warmup_steps	int	Step or timestep value.	1000
lr_warmup_start_ratio	float	Input value.	0.1
val_inference_sampler	str	Input value.	'ddpm'
val_ddim_num_timesteps	int	Input value.	200
val_ddim_eta	float	Input value.	0.0
log_intermediates	bool	Boolean flag controlling behavior.	True
skip_full_reconstruction_in_sanity_check	bool	Boolean flag controlling behavior.	True
max_full_reconstruction_samples	int	Input value.	4
postprocess_gaussian_blur_enabled	bool	Boolean flag controlling behavior.	False
postprocess_gaussian_blur_sigma	float	Input value.	0.35
postprocess_gaussian_blur_kernel_size	int	Input value.	3
wandb_verbose	bool	Boolean flag controlling behavior.	True
log_stats_every_n_steps	int	Step or timestep value.	1
log_images_every_n_steps	int	Step or timestep value.	200

Returns:

Name	Type	Description
None	None	No value is returned.

configure_optimizers()

Create optimizer and optional scheduler configuration.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Type	Description
Optimizer | dict[str, Any]	torch.optim.Optimizer | dict[str, Any]: Computed output value.

forward(condition, sampler=None, verbose=False, clamp_known_pixels=None, *, known_mask=None, known_values=None, coords=None, date=None, return_intermediates=False, intermediate_step_indices=None, return_x0_intermediates=False)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
condition	Tensor	Tensor input for the computation.	required
sampler	Module | None	Sampler instance used for reverse diffusion.	None
verbose	bool	Boolean flag controlling behavior.	False
clamp_known_pixels	bool | None	Boolean flag controlling behavior.	None
known_mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
known_values	Tensor | None	Tensor input for the computation.	None
coords	Tensor | None	Coordinate conditioning values.	None
date	Tensor | None	Date conditioning values.	None
return_intermediates	bool	Boolean flag controlling behavior.	False
intermediate_step_indices	list[int] | None	Input value.	None
return_x0_intermediates	bool	Boolean flag controlling behavior.	False

Returns:

Type	Description
Tensor | tuple[Tensor, list[tuple[int, Tensor]]] | tuple[Tensor, list[tuple[int, Tensor]], list[tuple[int, Tensor]]]	torch.Tensor | tuple[torch.Tensor, list[tuple[int, torch.Tensor]]] | tuple[torch.Tensor, list[tuple[int, torch.Tensor]], list[tuple[int, torch.Tensor]]]: Tensor output produced by this call.

from_config(model_config_path='configs/model_config.yaml', data_config_path='configs/data_config.yaml', training_config_path='configs/training_config.yaml', datamodule=None) classmethod

Compute from config and return the result.

Parameters:

Name	Type	Description	Default
model_config_path	str	Path to an input or output file.	'configs/model_config.yaml'
data_config_path	str	Path to an input or output file.	'configs/data_config.yaml'
training_config_path	str	Path to an input or output file.	'configs/training_config.yaml'
datamodule	LightningDataModule | None	Input value.	None

Returns:

Type	Description
'PixelDiffusionConditional'	'PixelDiffusionConditional': Computed output value.

input_T(value)

Compute input T and return the result.

Parameters:

Name	Type	Description	Default
value	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

load_state_dict(state_dict, strict=True)

Load checkpoint weights into the current module.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Tensor]	Tensor input for the computation.	required
strict	bool	Boolean flag controlling behavior.	True

Returns:

Name	Type	Description
Any	Any	Computed output value.

on_validation_epoch_end()

Compute on validation epoch end and return the result.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
None	None	No value is returned.

on_validation_epoch_start()

Compute on validation epoch start and return the result.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
None	None	No value is returned.

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)

Perform one optimizer step with optional learning-rate warmup.

Parameters:

Name	Type	Description	Default
epoch	int	Step or timestep value.	required
batch_idx	int	Zero-based index for selecting a sample or batch.	required
optimizer	Optimizer	Optimizer used for parameter updates.	required
optimizer_closure	Any | None	Optimizer used for parameter updates.	None

Returns:

Name	Type	Description
None	None	No value is returned.

output_T(value)

Compute output T and return the result.

Parameters:

Name	Type	Description	Default
value	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

predict_step(batch, batch_idx, dataloader_idx=0)

Compute predict step and return the result.

Parameters:

Name	Type	Description	Default
batch	dict[str, Any]	Input value.	required
batch_idx	int	Zero-based index for selecting a sample or batch.	required
dataloader_idx	int	Input value.	0

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Dictionary containing computed outputs.

train_dataloader()

Return the training dataloader from the attached datamodule.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Type	Description
DataLoader[Any]	torch.utils.data.DataLoader[Any]: Computed output value.

training_step(batch, batch_idx)

Compute training step and return the result.

Parameters:

Name	Type	Description	Default
batch	dict[str, Any]	Input value.	required
batch_idx	int	Zero-based index for selecting a sample or batch.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

val_dataloader()

Return the validation dataloader from the attached datamodule.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Type	Description
DataLoader[Any] | None	torch.utils.data.DataLoader[Any] | None: Computed output value.

validation_step(batch, batch_idx)

Compute validation step and return the result.

Parameters:

Name	Type	Description	Default
batch	dict[str, Any]	Input value.	required
batch_idx	int	Zero-based index for selecting a sample or batch.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

models.difFF.EMA

models.difFF.EMA

EMA

Bases: Callback

Callback that maintains exponential moving-average model weights.

ema_initialized property

Compute ema initialized and return the result.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Name	Type	Description
bool	bool	Computed scalar output.

__init__(decay, apply_ema_every_n_steps=1, start_step=0, save_ema_weights_in_callback_state=False, evaluate_ema_weights_instead=False)

Initialize EMA with configured parameters.

Parameters:

Name	Type	Description	Default
decay	float	Input value.	required
apply_ema_every_n_steps	int	Step or timestep value.	1
start_step	int	Step or timestep value.	0
save_ema_weights_in_callback_state	bool	Boolean flag controlling behavior.	False
evaluate_ema_weights_instead	bool	Boolean flag controlling behavior.	False

Returns:

Name	Type	Description
None	None	No value is returned.

apply_ema(pl_module)

Compute apply ema and return the result.

Parameters:

Name	Type	Description	Default
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

apply_multi_tensor_ema(pl_module)

Compute apply multi tensor ema and return the result.

Parameters:

Name	Type	Description	Default
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

ema(pl_module)

Compute ema and return the result.

Parameters:

Name	Type	Description	Default
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

load_state_dict(state_dict)

Load checkpoint weights into the current module.

Parameters:

Name	Type	Description	Default
state_dict	Dict[str, Any]	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

on_load_checkpoint(trainer, pl_module, checkpoint)

Compute on load checkpoint and return the result.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Input value.	required
pl_module	LightningModule	Input value.	required
checkpoint	Dict[str, Any]	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

on_test_end(trainer, pl_module)

Compute on test end and return the result.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Input value.	required
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

on_test_start(trainer, pl_module)

Compute on test start and return the result.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Input value.	required
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)

Compute on train batch end and return the result.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Input value.	required
pl_module	LightningModule	Input value.	required
outputs	STEP_OUTPUT	Input value.	required
batch	Any	Input value.	required
batch_idx	int	Zero-based index for selecting a sample or batch.	required

Returns:

Name	Type	Description
None	None	No value is returned.

on_train_start(trainer, pl_module)

Compute on train start and return the result.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Input value.	required
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

on_validation_end(trainer, pl_module)

Compute on validation end and return the result.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Input value.	required
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

on_validation_start(trainer, pl_module)

Compute on validation start and return the result.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Input value.	required
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

replace_model_weights(pl_module)

Compute replace model weights and return the result.

Parameters:

Name	Type	Description	Default
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

restore_original_weights(pl_module)

Compute restore original weights and return the result.

Parameters:

Name	Type	Description	Default
pl_module	LightningModule	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

should_apply_ema(step)

Compute should apply ema and return the result.

Parameters:

Name	Type	Description	Default
step	int	Step or timestep value.	required

Returns:

Name	Type	Description
bool	bool	Computed scalar output.

state_dict()

Return the serializable state for this object.

Parameters:

Name	Type	Description	Default
None		This callable takes no explicit input arguments.	required

Returns:

Type	Description
Dict[str, Any]	Dict[str, Any]: Computed output value.

models.difFF.DenoisingDiffusionProcess.DenoisingDiffusionProcess

models.difFF.DenoisingDiffusionProcess.DenoisingDiffusionProcess

ConvNextBlock

Bases: Module

ConvNeXt residual block used within the U-Net backbone.

__init__(dim, dim_out, *, time_emb_dim=None, coord_emb_dim=None, mult=2, norm=True)

Initialize ConvNextBlock with configured parameters.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required
dim_out	int	Input value.	required
time_emb_dim	int | None	Input value.	None
coord_emb_dim	int | None	Input value.	None
mult	int	Input value.	2
norm	bool	Boolean flag controlling behavior.	True

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x, time_emb=None, coord_emb=None)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor input for the computation.	required
time_emb	Tensor | None	Tensor input for the computation.	None
coord_emb	Tensor | None	Tensor input for the computation.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

DDIM_Sampler

Bases: Module

DDIM sampler that performs accelerated reverse-diffusion updates.

__init__(num_timesteps=100, train_timesteps=1000, clip_sample=True, schedule='linear', beta_start=0.0001, beta_end=0.02, eta=0.0, betas=None, parameterization='epsilon')

Initialize DDIM_Sampler with configured parameters.

Parameters:

Name	Type	Description	Default
num_timesteps	int	Step or timestep value.	100
train_timesteps	int	Step or timestep value.	1000
clip_sample	bool	Boolean flag controlling behavior.	True
schedule	str	Input value.	'linear'
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02
eta	float	Input value.	0.0
betas	Tensor | list[float] | tuple[float, ...] | None	Tensor input for the computation.	None
parameterization	str	Input value.	'epsilon'

Returns:

Name	Type	Description
None	None	No value is returned.

estimate_std(alpha_cumprod, alpha_cumprod_prev)

Compute estimate std and return the result.

Parameters:

Name	Type	Description	Default
alpha_cumprod	Tensor	Tensor input for the computation.	required
alpha_cumprod_prev	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

forward(*args, **kwargs)

Run the sampler call and return the next sample.

Parameters:

Name	Type	Description	Default
*args	Any	Additional positional arguments forwarded to the underlying call.	()
**kwargs	Any	Additional keyword arguments forwarded to the underlying call.	{}

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

set_parameterization(parameterization)

Compute set parameterization and return the result.

Parameters:

Name	Type	Description	Default
parameterization	str	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

step(x_t, t, z_t)

Predict the previous diffusion sample for one timestep.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
z_t	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

DDPM_Sampler

Bases: Module

DDPM sampler that performs one reverse-diffusion step at a time.

__init__(num_timesteps=1000, schedule='linear', beta_start=0.0001, beta_end=0.02, parameterization='epsilon')

Initialize DDPM_Sampler with configured parameters.

Parameters:

Name	Type	Description	Default
num_timesteps	int	Step or timestep value.	1000
schedule	str	Input value.	'linear'
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02
parameterization	str	Input value.	'epsilon'

Returns:

Name	Type	Description
None	None	No value is returned.

forward(*args, **kwargs)

Run the sampler call and return the next sample.

Parameters:

Name	Type	Description	Default
*args	Any	Additional positional arguments forwarded to the underlying call.	()
**kwargs	Any	Additional keyword arguments forwarded to the underlying call.	{}

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

posterior_params(x_t, t, noise_pred)

Compute posterior params and return the result.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
noise_pred	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
tuple[Tensor, Tensor]	tuple[torch.Tensor, torch.Tensor]: Tuple containing computed outputs.

set_parameterization(parameterization)

Compute set parameterization and return the result.

Parameters:

Name	Type	Description	Default
parameterization	str	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

step(x_t, t, z_t)

Predict the previous diffusion sample for one timestep.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
z_t	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

DenoisingDiffusionConditionalProcess

Bases: Module

Conditional diffusion process module for guided reconstruction.

__init__(generated_channels=3, condition_channels=3, loss_fn=F.mse_loss, schedule='linear', beta_start=0.0001, beta_end=0.02, num_timesteps=1000, unet_dim=64, unet_dim_mults=(1, 2, 4, 8), unet_with_time_emb=True, unet_output_mean_scale=False, unet_residual=False, coord_conditioning_enabled=False, coord_encoding='unit_sphere', date_conditioning_enabled=False, date_encoding='day_of_year_sincos', coord_embed_dim=None, parameterization='epsilon', sampler=None)

Initialize DenoisingDiffusionConditionalProcess with configured parameters.

Parameters:

Name	Type	Description	Default
generated_channels	int	Input value.	3
condition_channels	int	Input value.	3
loss_fn	Callable[[Tensor, Tensor], Tensor]	Tensor input for the computation.	mse_loss
schedule	str	Input value.	'linear'
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02
num_timesteps	int	Step or timestep value.	1000
unet_dim	int	Input value.	64
unet_dim_mults	tuple[int, ...]	Input value.	(1, 2, 4, 8)
unet_with_time_emb	bool	Boolean flag controlling behavior.	True
unet_output_mean_scale	bool	Boolean flag controlling behavior.	False
unet_residual	bool	Boolean flag controlling behavior.	False
coord_conditioning_enabled	bool	Boolean flag controlling behavior.	False
coord_encoding	str	Input value.	'unit_sphere'
date_conditioning_enabled	bool	Boolean flag controlling behavior.	False
date_encoding	str	Input value.	'day_of_year_sincos'
coord_embed_dim	int | None	Input value.	None
parameterization	str	Input value.	'epsilon'
sampler	Module | None	Sampler instance used for reverse diffusion.	None

Returns:

Name	Type	Description
None	None	No value is returned.

forward(condition, sampler=None, verbose=False, known_mask=None, known_values=None, coord=None, date=None, return_intermediates=False, intermediate_step_indices=None, return_x0_intermediates=False)

Run reverse diffusion and return generated outputs.

Parameters:

Name	Type	Description	Default
condition	Tensor	Tensor input for the computation.	required
sampler	Module | None	Sampler instance used for reverse diffusion.	None
verbose	bool	Boolean flag controlling behavior.	False
known_mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
known_values	Tensor | None	Tensor input for the computation.	None
coord	Tensor | None	Coordinate conditioning values.	None
date	Tensor | None	Date conditioning values.	None
return_intermediates	bool	Boolean flag controlling behavior.	False
intermediate_step_indices	list[int] | None	Input value.	None
return_x0_intermediates	bool	Boolean flag controlling behavior.	False

Returns:

Type	Description
Tensor | tuple[Tensor, list[tuple[int, Tensor]]] | tuple[Tensor, list[tuple[int, Tensor]], list[tuple[int, Tensor]]]	torch.Tensor | tuple[torch.Tensor, list[tuple[int, torch.Tensor]]] | tuple[torch.Tensor, list[tuple[int, torch.Tensor]], list[tuple[int, torch.Tensor]]]: Tensor output produced by this call.

p_loss(output, condition, *, valid_mask=None, land_mask=None, mask_loss=False, coord=None, date=None)

Compute the diffusion training loss for the current batch.

Parameters:

Name	Type	Description	Default
output	Tensor	Tensor input for the computation.	required
condition	Tensor	Tensor input for the computation.	required
valid_mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
land_mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
mask_loss	bool	Mask tensor controlling valid or known pixels.	False
coord	Tensor | None	Coordinate conditioning values.	None
date	Tensor | None	Date conditioning values.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

DenoisingDiffusionProcess

Bases: Module

Unconditional diffusion process module for training and sampling.

__init__(generated_channels=3, loss_fn=F.mse_loss, schedule='linear', beta_start=0.0001, beta_end=0.02, num_timesteps=1000, unet_dim=64, unet_dim_mults=(1, 2, 4, 8), unet_with_time_emb=True, unet_output_mean_scale=False, unet_residual=False, parameterization='epsilon', sampler=None)

Initialize DenoisingDiffusionProcess with configured parameters.

Parameters:

Name	Type	Description	Default
generated_channels	int	Input value.	3
loss_fn	Callable[[Tensor, Tensor], Tensor]	Tensor input for the computation.	mse_loss
schedule	str	Input value.	'linear'
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02
num_timesteps	int	Step or timestep value.	1000
unet_dim	int	Input value.	64
unet_dim_mults	tuple[int, ...]	Input value.	(1, 2, 4, 8)
unet_with_time_emb	bool	Boolean flag controlling behavior.	True
unet_output_mean_scale	bool	Boolean flag controlling behavior.	False
unet_residual	bool	Boolean flag controlling behavior.	False
parameterization	str	Input value.	'epsilon'
sampler	Module | None	Sampler instance used for reverse diffusion.	None

Returns:

Name	Type	Description
None	None	No value is returned.

forward(shape=(256, 256), batch_size=1, sampler=None, verbose=False)

Run reverse diffusion and return generated outputs.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	Input value.	(256, 256)
batch_size	int	Size/count parameter.	1
sampler	Module | None	Sampler instance used for reverse diffusion.	None
verbose	bool	Boolean flag controlling behavior.	False

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

p_loss(output)

Compute the diffusion training loss for the current batch.

Parameters:

Name	Type	Description	Default
output	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

ForwardModel

Bases: Module

Base interface for forward diffusion process implementations.

__init__(num_timesteps=1000, schedule='linear')

Initialize ForwardModel with configured parameters.

Parameters:

Name	Type	Description	Default
num_timesteps	int	Step or timestep value.	1000
schedule	str	Input value.	'linear'

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x_0, t)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x_0	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

step(x_t, t)

Run one update step.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

GaussianForwardProcess

Bases: ForwardModel

Forward diffusion process based on Gaussian noise transitions.

__init__(num_timesteps=1000, schedule='linear', beta_start=0.0001, beta_end=0.02)

Initialize GaussianForwardProcess with configured parameters.

Parameters:

Name	Type	Description	Default
num_timesteps	int	Step or timestep value.	1000
schedule	str	Input value.	'linear'
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x_0, t, return_noise=False)

Run reverse diffusion and return generated outputs.

Parameters:

Name	Type	Description	Default
x_0	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
return_noise	bool	Boolean flag controlling behavior.	False

Returns:

Type	Description
Tensor | tuple[Tensor, Tensor]	torch.Tensor | tuple[torch.Tensor, torch.Tensor]: Tensor output produced by this call.

step(x_t, t, return_noise=False)

Apply one forward-diffusion transition step.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
return_noise	bool	Boolean flag controlling behavior.	False

Returns:

Type	Description
Tensor | tuple[Tensor, Tensor]	torch.Tensor | tuple[torch.Tensor, torch.Tensor]: Tensor output produced by this call.

LayerNorm

Bases: Module

Channel-wise layer normalization for 2D feature maps.

__init__(dim, eps=1e-05)

Initialize LayerNorm with configured parameters.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required
eps	float	Input value.	1e-05

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

LinearAttention

Bases: Module

Linear attention block for efficient spatial mixing.

__init__(dim, heads=4, dim_head=32)

Initialize LinearAttention with configured parameters.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required
heads	int	Input value.	4
dim_head	int	Input value.	32

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

PreNorm

Bases: Module

Module that normalizes inputs before applying a submodule.

__init__(dim, fn)

Initialize PreNorm with configured parameters.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required
fn	Module	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

Residual

Bases: Module

Wrapper module that adds a residual skip connection.

__init__(fn)

Initialize Residual with configured parameters.

Parameters:

Name	Type	Description	Default
fn	Module	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x, *args, **kwargs)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor input for the computation.	required
*args	Any	Additional positional arguments forwarded to the underlying call.	()
**kwargs	Any	Additional keyword arguments forwarded to the underlying call.	{}

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

SinusoidalPosEmb

Bases: Module

Module that generates sinusoidal timestep embeddings.

__init__(dim)

Initialize SinusoidalPosEmb with configured parameters.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

UnetConvNextBlock

Bases: Module

U-Net/ConvNeXt backbone used by the diffusion model.

__init__(dim, out_dim=None, dim_mults=(1, 2, 4, 8), channels=3, with_time_emb=True, coord_emb_dim=None, output_mean_scale=False, residual=False)

Initialize UnetConvNextBlock with configured parameters.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required
out_dim	int | None	Input value.	None
dim_mults	tuple[int, ...]	Input value.	(1, 2, 4, 8)
channels	int	Input value.	3
with_time_emb	bool	Boolean flag controlling behavior.	True
coord_emb_dim	int | None	Input value.	None
output_mean_scale	bool	Boolean flag controlling behavior.	False
residual	bool	Boolean flag controlling behavior.	False

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x, time=None, coord_emb=None)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor input for the computation.	required
time	Tensor | None	Tensor input for the computation.	None
coord_emb	Tensor | None	Tensor input for the computation.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

Downsample(dim)

Create a strided-convolution downsampling layer.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required

Returns:

Type	Description
Conv2d	nn.Conv2d: Computed output value.

Upsample(dim)

Create a transpose-convolution upsampling layer.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required

Returns:

Type	Description
ConvTranspose2d	nn.ConvTranspose2d: Computed output value.

cosine_beta_schedule(timesteps, s=0.008, beta_start=0.0001, beta_end=None)

Compute cosine beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
s	float	Input value.	0.008
beta_start	float | None	Input value.	0.0001
beta_end	float | None	Input value.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

default(val, d)

Return the input value or a fallback default.

Parameters:

Name	Type	Description	Default
val	T | None	Input value.	required
d	T | Callable[[], T]	Input value.	required

Returns:

Name	Type	Description
T	T	Computed output value.

exists(x)

Return whether the provided value is not None.

Parameters:

Name	Type	Description	Default
x	object	Input value.	required

Returns:

Name	Type	Description
bool	bool	Computed scalar output.

get_beta_schedule(variant, timesteps, beta_start=0.0001, beta_end=0.02)

Compute get beta schedule and return the result.

Parameters:

Name	Type	Description	Default
variant	str	Input value.	required
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute linear beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

quadratic_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute quadratic beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

sigmoid_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute sigmoid beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

models.difFF.DenoisingDiffusionProcess.forward

models.difFF.DenoisingDiffusionProcess.forward

This file contains implementations of the forward diffusion process

Current Models:

1) Gaussian Diffusion

ForwardModel

Bases: Module

Base interface for forward diffusion process implementations.

__init__(num_timesteps=1000, schedule='linear')

Initialize ForwardModel with configured parameters.

Parameters:

Name	Type	Description	Default
num_timesteps	int	Step or timestep value.	1000
schedule	str	Input value.	'linear'

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x_0, t)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x_0	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

step(x_t, t)

Run one update step.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

GaussianForwardProcess

Bases: ForwardModel

Forward diffusion process based on Gaussian noise transitions.

__init__(num_timesteps=1000, schedule='linear', beta_start=0.0001, beta_end=0.02)

Initialize GaussianForwardProcess with configured parameters.

Parameters:

Name	Type	Description	Default
num_timesteps	int	Step or timestep value.	1000
schedule	str	Input value.	'linear'
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x_0, t, return_noise=False)

Run reverse diffusion and return generated outputs.

Parameters:

Name	Type	Description	Default
x_0	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
return_noise	bool	Boolean flag controlling behavior.	False

Returns:

Type	Description
Tensor | tuple[Tensor, Tensor]	torch.Tensor | tuple[torch.Tensor, torch.Tensor]: Tensor output produced by this call.

step(x_t, t, return_noise=False)

Apply one forward-diffusion transition step.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
return_noise	bool	Boolean flag controlling behavior.	False

Returns:

Type	Description
Tensor | tuple[Tensor, Tensor]	torch.Tensor | tuple[torch.Tensor, torch.Tensor]: Tensor output produced by this call.

cosine_beta_schedule(timesteps, s=0.008, beta_start=0.0001, beta_end=None)

Compute cosine beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
s	float	Input value.	0.008
beta_start	float | None	Input value.	0.0001
beta_end	float | None	Input value.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

get_beta_schedule(variant, timesteps, beta_start=0.0001, beta_end=0.02)

Compute get beta schedule and return the result.

Parameters:

Name	Type	Description	Default
variant	str	Input value.	required
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute linear beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

quadratic_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute quadratic beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

sigmoid_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute sigmoid beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

models.difFF.DenoisingDiffusionProcess.beta_schedules

models.difFF.DenoisingDiffusionProcess.beta_schedules

cosine_beta_schedule(timesteps, s=0.008, beta_start=0.0001, beta_end=None)

Compute cosine beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
s	float	Input value.	0.008
beta_start	float | None	Input value.	0.0001
beta_end	float | None	Input value.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

get_beta_schedule(variant, timesteps, beta_start=0.0001, beta_end=0.02)

Compute get beta schedule and return the result.

Parameters:

Name	Type	Description	Default
variant	str	Input value.	required
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute linear beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

quadratic_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute quadratic beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

sigmoid_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute sigmoid beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

models.difFF.DenoisingDiffusionProcess.samplers.DDPM

models.difFF.DenoisingDiffusionProcess.samplers.DDPM

This file contains the DDPM sampler class for a diffusion process

DDPM_Sampler

Bases: Module

DDPM sampler that performs one reverse-diffusion step at a time.

__init__(num_timesteps=1000, schedule='linear', beta_start=0.0001, beta_end=0.02, parameterization='epsilon')

Initialize DDPM_Sampler with configured parameters.

Parameters:

Name	Type	Description	Default
num_timesteps	int	Step or timestep value.	1000
schedule	str	Input value.	'linear'
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02
parameterization	str	Input value.	'epsilon'

Returns:

Name	Type	Description
None	None	No value is returned.

forward(*args, **kwargs)

Run the sampler call and return the next sample.

Parameters:

Name	Type	Description	Default
*args	Any	Additional positional arguments forwarded to the underlying call.	()
**kwargs	Any	Additional keyword arguments forwarded to the underlying call.	{}

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

posterior_params(x_t, t, noise_pred)

Compute posterior params and return the result.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
noise_pred	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
tuple[Tensor, Tensor]	tuple[torch.Tensor, torch.Tensor]: Tuple containing computed outputs.

set_parameterization(parameterization)

Compute set parameterization and return the result.

Parameters:

Name	Type	Description	Default
parameterization	str	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

step(x_t, t, z_t)

Predict the previous diffusion sample for one timestep.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
z_t	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

cosine_beta_schedule(timesteps, s=0.008, beta_start=0.0001, beta_end=None)

Compute cosine beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
s	float	Input value.	0.008
beta_start	float | None	Input value.	0.0001
beta_end	float | None	Input value.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

get_beta_schedule(variant, timesteps, beta_start=0.0001, beta_end=0.02)

Compute get beta schedule and return the result.

Parameters:

Name	Type	Description	Default
variant	str	Input value.	required
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute linear beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

quadratic_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute quadratic beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

sigmoid_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute sigmoid beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

models.difFF.DenoisingDiffusionProcess.samplers.DDIM

models.difFF.DenoisingDiffusionProcess.samplers.DDIM

This file contains the DDIM sampler class for a diffusion process

DDIM_Sampler

Bases: Module

DDIM sampler that performs accelerated reverse-diffusion updates.

__init__(num_timesteps=100, train_timesteps=1000, clip_sample=True, schedule='linear', beta_start=0.0001, beta_end=0.02, eta=0.0, betas=None, parameterization='epsilon')

Initialize DDIM_Sampler with configured parameters.

Parameters:

Name	Type	Description	Default
num_timesteps	int	Step or timestep value.	100
train_timesteps	int	Step or timestep value.	1000
clip_sample	bool	Boolean flag controlling behavior.	True
schedule	str	Input value.	'linear'
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02
eta	float	Input value.	0.0
betas	Tensor | list[float] | tuple[float, ...] | None	Tensor input for the computation.	None
parameterization	str	Input value.	'epsilon'

Returns:

Name	Type	Description
None	None	No value is returned.

estimate_std(alpha_cumprod, alpha_cumprod_prev)

Compute estimate std and return the result.

Parameters:

Name	Type	Description	Default
alpha_cumprod	Tensor	Tensor input for the computation.	required
alpha_cumprod_prev	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

forward(*args, **kwargs)

Run the sampler call and return the next sample.

Parameters:

Name	Type	Description	Default
*args	Any	Additional positional arguments forwarded to the underlying call.	()
**kwargs	Any	Additional keyword arguments forwarded to the underlying call.	{}

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

set_parameterization(parameterization)

Compute set parameterization and return the result.

Parameters:

Name	Type	Description	Default
parameterization	str	Input value.	required

Returns:

Name	Type	Description
None	None	No value is returned.

step(x_t, t, z_t)

Predict the previous diffusion sample for one timestep.

Parameters:

Name	Type	Description	Default
x_t	Tensor	Tensor input for the computation.	required
t	Tensor	Tensor input for the computation.	required
z_t	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

cosine_beta_schedule(timesteps, s=0.008, beta_start=0.0001, beta_end=None)

Compute cosine beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
s	float	Input value.	0.008
beta_start	float | None	Input value.	0.0001
beta_end	float | None	Input value.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

get_beta_schedule(variant, timesteps, beta_start=0.0001, beta_end=0.02)

Compute get beta schedule and return the result.

Parameters:

Name	Type	Description	Default
variant	str	Input value.	required
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute linear beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

quadratic_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute quadratic beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

sigmoid_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02)

Compute sigmoid beta schedule and return the result.

Parameters:

Name	Type	Description	Default
timesteps	int	Step or timestep value.	required
beta_start	float	Input value.	0.0001
beta_end	float	Input value.	0.02

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

models.difFF.DenoisingDiffusionProcess.DenoisingDiffusionProcess.UnetConvNextBlock

models.difFF.DenoisingDiffusionProcess.DenoisingDiffusionProcess.UnetConvNextBlock

Bases: Module

U-Net/ConvNeXt backbone used by the diffusion model.

__init__(dim, out_dim=None, dim_mults=(1, 2, 4, 8), channels=3, with_time_emb=True, coord_emb_dim=None, output_mean_scale=False, residual=False)

Initialize UnetConvNextBlock with configured parameters.

Parameters:

Name	Type	Description	Default
dim	int	Input value.	required
out_dim	int | None	Input value.	None
dim_mults	tuple[int, ...]	Input value.	(1, 2, 4, 8)
channels	int	Input value.	3
with_time_emb	bool	Boolean flag controlling behavior.	True
coord_emb_dim	int | None	Input value.	None
output_mean_scale	bool	Boolean flag controlling behavior.	False
residual	bool	Boolean flag controlling behavior.	False

Returns:

Name	Type	Description
None	None	No value is returned.

forward(x, time=None, coord_emb=None)

Run the module forward computation.

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor input for the computation.	required
time	Tensor | None	Tensor input for the computation.	None
coord_emb	Tensor | None	Tensor input for the computation.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

Utilities

utils.normalizations

utils.normalizations

temperature_normalize(mode, tensor)

Compute temperature normalize and return the result.

Parameters:

Name	Type	Description	Default
mode	str	Input value.	required
tensor	Tensor	Tensor input for the computation.	required

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

temperature_to_plot_unit(tensor, *, tensor_is_normalized=True)

Compute temperature to plot unit and return the result.

Parameters:

Name	Type	Description	Default
tensor	Tensor	Tensor input for the computation.	required
tensor_is_normalized	bool	Boolean flag controlling behavior.	True

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

utils.stretching

utils.stretching

minmax_stretch(tensor, *, mask=None, nodata_value=None)

Compute minmax stretch and return the result.

Parameters:

Name	Type	Description	Default
tensor	Tensor	Tensor input for the computation.	required
mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
nodata_value	float | None	Input value.	None

Returns:

Type	Description
Tensor	torch.Tensor: Tensor output produced by this call.

utils.validation_denoise

utils.validation_denoise

build_capture_indices(total_steps, intermediate_step_indices)

Build and return capture indices.

Parameters:

Name	Type	Description	Default
total_steps	int	Step or timestep value.	required
intermediate_step_indices	list[int] | None	Input value.	required

Returns:

Type	Description
set[int]	set[int]: Computed output value.

build_evenly_spaced_capture_steps(total_steps, num_frames)

Build and return evenly spaced capture steps.

Parameters:

Name	Type	Description	Default
total_steps	int	Step or timestep value.	required
num_frames	int	Size/count parameter.	required

Returns:

Type	Description
list[int]	list[int]: List containing computed outputs.

log_wandb_conditional_reconstruction_grid(*, logger, x, y_hat, y_target, valid_mask=None, land_mask=None, eo=None, prefix='val_imgs', image_key='x_y_full_reconstruction', cmap='turbo')

Log wandb conditional reconstruction grid for monitoring.

Parameters:

Name	Type	Description	Default
logger	Any	Logger instance used for experiment tracking.	required
x	Tensor	Tensor input for the computation.	required
y_hat	Tensor	Tensor input for the computation.	required
y_target	Tensor	Tensor input for the computation.	required
valid_mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
land_mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
eo	Tensor | None	Tensor input for the computation.	None
prefix	str	Input value.	'val_imgs'
image_key	str	Input value.	'x_y_full_reconstruction'
cmap	str	Input value.	'turbo'

Returns:

Name	Type	Description
None	None	No value is returned.

log_wandb_denoise_timestep_grid(*, logger, denoise_samples, mae_samples=None, total_steps, sampler, conditioning_image=None, eo_conditioning_image=None, ground_truth=None, valid_mask=None, land_mask=None, prefix='val_imgs', cmap='turbo', nrows=4, ncols=4, tile_size_px=128, tile_pad_px=2)

Log wandb denoise timestep grid for monitoring.

Parameters:

Name	Type	Description	Default
logger	Any	Logger instance used for experiment tracking.	required
denoise_samples	list[tuple[int, Tensor]]	Tensor input for the computation.	required
mae_samples	list[tuple[int, Tensor]] | None	Tensor input for the computation.	None
total_steps	int	Step or timestep value.	required
sampler	Any	Sampler instance used for reverse diffusion.	required
conditioning_image	Tensor | None	Tensor input for the computation.	None
eo_conditioning_image	Tensor | None	Tensor input for the computation.	None
ground_truth	Tensor | None	Tensor input for the computation.	None
valid_mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
land_mask	Tensor | None	Mask tensor controlling valid or known pixels.	None
prefix	str	Input value.	'val_imgs'
cmap	str	Input value.	'turbo'
nrows	int	Input value.	4
ncols	int	Input value.	4
tile_size_px	int	Input value.	128
tile_pad_px	int	Input value.	2

Returns:

Name	Type	Description
None	None	No value is returned.

log_wandb_diffusion_schedule_profile(*, logger, sampler, total_steps, prefix='val_imgs', eps=1e-12)

Log wandb diffusion schedule profile for monitoring.

Parameters:

Name	Type	Description	Default
logger	Any	Logger instance used for experiment tracking.	required
sampler	Any	Sampler instance used for reverse diffusion.	required
total_steps	int	Step or timestep value.	required
prefix	str	Input value.	'val_imgs'
eps	float	Input value.	1e-12

Returns:

Name	Type	Description
None	None	No value is returned.

step_to_sampler_timestep_label(*, step_index, total_steps, sampler)

Compute step to sampler timestep label and return the result.

Parameters:

Name	Type	Description	Default
step_index	int	Input value.	required
total_steps	int	Step or timestep value.	required
sampler	Any	Sampler instance used for reverse diffusion.	required

Returns:

Name	Type	Description
int	int	Computed scalar output.