diff --git a/README.md b/README.md index 934050f..e0e5ce7 100644 --- a/README.md +++ b/README.md @@ -1,19 +1,55 @@ -# cosmo_diffusion +![banner](docs/_static/banner.png) -Train 2D/3D diffusion models (UNet and DiT) for cosmological applications. +Train 2D/3D diffusion (and flow-matching) models — UNet, UNet-Conditional, DiT, +and PixArt — for cosmological applications. ## Install + ```bash git clone https://github.com/nkern/cosmo_diffusion cd cosmo_diffusion -pip install . +pip install -e . ``` -## Running -Look at `cosmodiff/configs/config.yaml` for a configuration file. Then just run: +## Dependencies + +- `numpy` +- `torch` +- `diffusers` +- `accelerate` +- `tqdm` +- `pyyaml` +- `scipy` +- `h5py` +- `matplotlib` +- `ema-pytorch` + +## Quick demo + +Configure a training run in `cosmodiff/data/config.yaml` (paths, model, +scheduler, training kwargs), then launch: ```bash -cosmodiff_train.py --config path_to_config +cosmodiff_train.py --config path/to/config.yaml ``` -checkpointing and metrics are automatically stored in `output_dir` (defined in the config). +Checkpoints and metrics are written automatically to the `output_dir` set in +the config. To sample from a trained checkpoint: + +```bash +cosmodiff_sample.py --output_dir path/to/run \ + --n_samples 64 --output samples.npy +``` + +For fast inference, swap in a higher-order solver: + +```bash +cosmodiff_sample.py --output_dir path/to/run \ + --scheduler DPMSolverMultistepScheduler --num_steps 25 \ + --n_samples 64 --output samples.npy +``` + +## Authors + +- Nicholas Kern +- Jiaming Pan diff --git a/cosmodiff/__init__.py b/cosmodiff/__init__.py index 745e8bb..32a343b 100644 --- a/cosmodiff/__init__.py +++ b/cosmodiff/__init__.py @@ -1,6 +1,9 @@ +from importlib.metadata import version as _pkg_version + from . import utils from . import augment from . import optim from .optim import train, generate -from .version import __version__ + +__version__ = _pkg_version("cosmodiff") diff --git a/cosmodiff/augment.py b/cosmodiff/augment.py index a11ed93..f2bb9de 100644 --- a/cosmodiff/augment.py +++ b/cosmodiff/augment.py @@ -9,16 +9,17 @@ class RandomRoll(nn.Module): Args: dims (tuple of int): Dimensions to roll along. Defaults to ``(-1,)``. """ - def __init__(self, dims=(-1,)): + def __init__(self, size=128, dims=(-1,)): super().__init__() + self.size = size self.dims = tuple(dims) self.ndim = len(dims) def __call__(self, x): if x is None: return None - shift = (torch.rand(self.ndim, device='cpu') * torch.tensor(x.shape)[list(self.dims)]).round().to(torch.long) - return torch.roll(x, tuple(shift), self.dims) + shift = torch.randint(0, self.size, (self.ndim,), dtype=torch.long, device='cpu').tolist() + return torch.roll(x, shift, self.dims) def __repr__(self) -> str: return f"{self.__class__.__name__}(dims={self.dims})" @@ -68,15 +69,18 @@ class RandomFlip(nn.Module): """ def __init__(self, dims=(-2, -1), p=0.5): super().__init__() - self.dims = tuple(dims) + self.dims = list(dims) self.p = p def __call__(self, x): if x is None: return None - flip = torch.rand(len(self.dims), device='cpu') - flip = torch.where(flip < self.p)[0].tolist() - return torch.flip(x, flip) + draw = torch.rand(len(self.dims), device='cpu') + flip = torch.where(draw < self.p)[0].tolist() + if len(flip) == 0: + return x + else: + return torch.flip(x, [self.dims[f] for f in flip]) def __repr__(self) -> str: return f"{self.__class__.__name__}(dims={self.dims}, p={self.p})" @@ -119,4 +123,3 @@ def config_augmentations(augmentations): return nn.Sequential(*pipeline) - diff --git a/cosmodiff/configs/config.yaml b/cosmodiff/configs/config.yaml deleted file mode 100644 index 2161c35..0000000 --- a/cosmodiff/configs/config.yaml +++ /dev/null @@ -1,85 +0,0 @@ -# cosmodiff training config - -# --- global --- -global: - device: cpu - dtype: float32 - -# --- io --- -io: - output_dir: /path/to/output - -# --- data --- -data: - img_path: /path/to/images.npy - img_read_fn: npy_read_fn # any function in cosmodiff.utils - label_path: /path/to/labels.npy # optional - label_read_fn: npy_read_fn # optional - log: false - norm: center-scale - two_dim: true - zthin: 4 - n_samples: null - seed: null - keep_on_cpu: false - -# --- augmentations --- -augmentations: - RandomRoll: - dims: [-1, -2] - RandomFlip: - dims: [-1, -2] - p: 0.5 - -# --- model --- -model: - class: UNet2DModel - kwargs: - sample_size: 64 - in_channels: 1 - out_channels: 1 - layers_per_block: 2 - block_out_channels: [64, 128, 256] - down_block_types: - - DownBlock2D - - DownBlock2D - - DownBlock2D - up_block_types: - - UpBlock2D - - UpBlock2D - - UpBlock2D - norm_num_groups: 32 - -# --- noise scheduler --- -noise_scheduler: - class: DDPMScheduler - kwargs: - num_train_timesteps: 1000 - -# --- optimizer --- -optimizer: - class: AdamW - kwargs: - lr: 1.0e-4 - weight_decay: 1.0e-2 - -# --- lr scheduler --- -lr_scheduler: - class: ConstantLR - kwargs: - factor: 1.0 - total_iters: 0 - -# --- training --- -train: - num_epochs: 50 - batch_size: 16 - shuffle: true - checkpoint_every_n_epochs: 5 - mixed_precision: fp16 - gradient_accumulation_steps: 1 - dataloader_num_workers: 4 - max_grad_norm: 1.0 - verbose: true - force_cpu: false - pin_memory: false diff --git a/cosmodiff/data/IllustrisTNG_Mcdm.npy b/cosmodiff/data/IllustrisTNG_Mcdm.npy new file mode 100644 index 0000000..f83d933 Binary files /dev/null and b/cosmodiff/data/IllustrisTNG_Mcdm.npy differ diff --git a/cosmodiff/data/__init__.py b/cosmodiff/data/__init__.py new file mode 100644 index 0000000..ae4da01 --- /dev/null +++ b/cosmodiff/data/__init__.py @@ -0,0 +1,3 @@ +from pathlib import Path + +DATA_PATH = Path(__file__).parent diff --git a/cosmodiff/data/config.yaml b/cosmodiff/data/config.yaml new file mode 100644 index 0000000..2d69629 --- /dev/null +++ b/cosmodiff/data/config.yaml @@ -0,0 +1,119 @@ +# cosmodiff training config + +# --- global --- +global: + device: cpu + dtype: float32 + +# --- io --- +io: + output_dir: /path/to/output + +# --- data --- +data: + img_path: /path/to/images.npy + img_read_fn: npy_read_fn # any function in cosmodiff.utils + label_path: /path/to/labels.npy # optional + label_read_fn: npy_read_fn # optional + reshape: 2d # '2d' (z-slice → channel), '3d' (volume), or null (no reshape) + zthin: 4 # thinning factor along z when reshape='2d' + n_samples: null + seed: null + keep_on_cpu: true + normalization: center-max # data norm and kwargs + norm_kwargs: + center: null + xmax: null + alpha: null + beta: null + transform: + [log] + +# --- augmentations --- +augmentations: + RandomRoll: + size: 128 + dims: [-1, -2] + RandomFlip: + dims: [-1, -2] + p: 0.5 + +# --- model --- +model: + class: UNet2DModel + kwargs: + sample_size: 64 + in_channels: 1 + out_channels: 1 + layers_per_block: 2 + block_out_channels: [64, 128, 256] + down_block_types: + - DownBlock2D + - DownBlock2D + - DownBlock2D + up_block_types: + - UpBlock2D + - UpBlock2D + - UpBlock2D + norm_num_groups: 32 + +# --- noise scheduler --- +noise_scheduler: + class: DDPMScheduler + kwargs: + num_train_timesteps: 1000 + rescale_betas_zero_snr: true + +# --- optimizer --- +optimizer: + class: AdamW + kwargs: + lr: 1.0e-4 + weight_decay: 1.0e-2 + +# --- lr scheduler --- +lr_scheduler: + class: ConstantLR + kwargs: + factor: 1.0 + total_iters: 0 + +# --- training --- +train: + num_epochs: 50 + batch_size: 16 + shuffle: true + checkpoint_every_n_epochs: 5 + mixed_precision: fp16 + gradient_accumulation_steps: 1 + dataloader_num_workers: 4 + max_grad_norm: 1.0 + conditioning: discrete # 'discrete' (class_labels) or 'continuous' (encoder_hidden_states) + cfg_dropout: 0.0 # fraction of labels dropped for CFG training (0 disables) + ema_sigma_rels: [.02, .10] # e.g. [0.02, 0.10] to enable post-hoc EMA tracking + ema_update_every: 1 # optimizer steps between EMA updates (Karras assumes per-step) + ema_burn_in: 1000 # skip first N optimizer steps before starting EMA tracking + min_snr_gamma: null # set to e.g. 5.0 to enable Min-SNR loss weighting (Hang et al. 2023) + sigma_log_normal: null # set to e.g. [-1.2, 1.2] for EDM-style log-σ ~ Normal(P_mean, P_std) sampling + verbose: true + force_cpu: false + pin_memory: false + +# --- generation --- +generate: + scheduler: null # new scheduler for inference, default is train scheduler + num_steps: null # number of inference steps, default is training steps + s_churn: null # EDM stochasticity (Euler/Heun-family only); 0=ODE, larger=more SDE-like + s_tmin: null # restrict churn to t >= s_tmin + s_tmax: null # restrict churn to t <= s_tmax + s_noise: null # noise magnitude multiplier during churn (default 1.0) + n_samples: 64 # total number of samples to generate + batch_size: null # samples per forward pass; null → single batch of n_samples + image_shape: null # e.g. [1, 64, 64]; null → inferred from model.config + conditioning: discrete # 'discrete' or 'continuous' (must match training) + labels: null # discrete: list of int class labels (length 1 or n_samples) + continuous_labels: null # continuous: path to .npy of shape (n_samples, D) or (1, D) + guidance_scale: null # CFG guidance scale; null disables amplification + ema_sigma_rel: null # synthesize EMA at this target before sampling; null disables + seed: null # set for reproducibility + device: null # null → cuda if available else cpu diff --git a/cosmodiff/data/config_FM.yaml b/cosmodiff/data/config_FM.yaml new file mode 100644 index 0000000..fc808ac --- /dev/null +++ b/cosmodiff/data/config_FM.yaml @@ -0,0 +1,119 @@ +# cosmodiff training config — flow-matching variant + +# --- global --- +global: + device: cpu + dtype: float32 + +# --- io --- +io: + output_dir: /path/to/output + +# --- data --- +data: + img_path: /path/to/images.npy + img_read_fn: npy_read_fn # any function in cosmodiff.utils + label_path: /path/to/labels.npy # optional + label_read_fn: npy_read_fn # optional + reshape: 2d # '2d' (z-slice → channel), '3d' (volume), or null (no reshape) + zthin: 4 # thinning factor along z when reshape='2d' + n_samples: null + seed: null + keep_on_cpu: true + normalization: center-max # data norm and kwargs + norm_kwargs: + center: null + xmax: null + alpha: null + beta: null + transform: + [log] + +# --- augmentations --- +augmentations: + RandomRoll: + size: 128 + dims: [-1, -2] + RandomFlip: + dims: [-1, -2] + p: 0.5 + +# --- model --- +model: + class: UNet2DModel + kwargs: + sample_size: 64 + in_channels: 1 + out_channels: 1 + layers_per_block: 2 + block_out_channels: [64, 128, 256] + down_block_types: + - DownBlock2D + - DownBlock2D + - DownBlock2D + up_block_types: + - UpBlock2D + - UpBlock2D + - UpBlock2D + norm_num_groups: 32 + +# --- noise scheduler --- +noise_scheduler: + class: FlowMatchEulerDiscreteScheduler + kwargs: + num_train_timesteps: 1000 + # shift: 1.0 # optional time-axis warp; SD3 uses larger values (~3.0) + +# --- optimizer --- +optimizer: + class: AdamW + kwargs: + lr: 1.0e-4 + weight_decay: 1.0e-2 + +# --- lr scheduler --- +lr_scheduler: + class: ConstantLR + kwargs: + factor: 1.0 + total_iters: 0 + +# --- training --- +train: + num_epochs: 50 + batch_size: 16 + shuffle: true + checkpoint_every_n_epochs: 5 + mixed_precision: fp16 + gradient_accumulation_steps: 1 + dataloader_num_workers: 4 + max_grad_norm: 1.0 + conditioning: discrete # 'discrete' (class_labels) or 'continuous' (encoder_hidden_states) + cfg_dropout: 0.0 # fraction of labels dropped for CFG training (0 disables) + ema_sigma_rels: [.02, .10] # e.g. [0.02, 0.10] to enable post-hoc EMA tracking + ema_update_every: 1 # optimizer steps between EMA updates (Karras assumes per-step) + ema_burn_in: 1000 # skip first N optimizer steps before starting EMA tracking + min_snr_gamma: null # ignored for FM (DDPM-only); kept for config compatibility + sigma_log_normal: null # for FM, interpreted as (P_mean, P_std) of logit-normal `t` sampling (SD3-style) + verbose: true + force_cpu: false + pin_memory: false + +# --- generation --- +generate: + scheduler: null # e.g. FlowMatchHeunDiscreteScheduler for 2nd-order ODE; null → use training scheduler + num_steps: null # number of inference steps, default is training steps; FM works at 20-50 typically + s_churn: null # EDM stochasticity (Euler/Heun-family only); 0=ODE, larger=more SDE-like + s_tmin: null # restrict churn to t >= s_tmin + s_tmax: null # restrict churn to t <= s_tmax + s_noise: null # noise magnitude multiplier during churn (default 1.0) + n_samples: 64 # total number of samples to generate + batch_size: null # samples per forward pass; null → single batch of n_samples + image_shape: null # e.g. [1, 64, 64]; null → inferred from model.config + conditioning: discrete # 'discrete' or 'continuous' (must match training) + labels: null # discrete: list of int class labels (length 1 or n_samples) + continuous_labels: null # continuous: path to .npy of shape (n_samples, D) or (1, D) + guidance_scale: null # CFG guidance scale; null disables amplification + ema_sigma_rel: null # synthesize EMA at this target before sampling; null disables + seed: null # set for reproducibility + device: null # null → cuda if available else cpu diff --git a/cosmodiff/data/config_multipath.yaml b/cosmodiff/data/config_multipath.yaml new file mode 100644 index 0000000..2e1e87d --- /dev/null +++ b/cosmodiff/data/config_multipath.yaml @@ -0,0 +1,123 @@ +# cosmodiff training config + +# --- global --- +global: + device: cpu + dtype: float32 + +# --- io --- +io: + output_dir: /path/to/output + +# --- data --- +data: + img_path: + - /path/to/images.npy + - /path/to/images2.npy + img_read_fn: npy_read_fn + label_path: + - /path/to/labels.npy + - /path/to/labels.npy + label_read_fn: npy_read_fn + two_dim: true + zthin: 4 + n_samples: null + seed: null + keep_on_cpu: true + normalization: center-max # data norm and kwargs + norm_kwargs: + center: null + xmax: null + alpha: null + beta: null + transform: + [log] + +# --- augmentations --- +augmentations: + RandomRoll: + size: 128 + dims: [-1, -2] + RandomFlip: + dims: [-1, -2] + p: 0.5 + +# --- model --- +model: + class: UNet2DModel + kwargs: + sample_size: 64 + in_channels: 1 + out_channels: 1 + layers_per_block: 2 + block_out_channels: [64, 128, 256] + down_block_types: + - DownBlock2D + - DownBlock2D + - DownBlock2D + up_block_types: + - UpBlock2D + - UpBlock2D + - UpBlock2D + norm_num_groups: 32 + +# --- noise scheduler --- +noise_scheduler: + class: DDPMScheduler + kwargs: + num_train_timesteps: 1000 + rescale_betas_zero_snr: true + +# --- optimizer --- +optimizer: + class: AdamW + kwargs: + lr: 1.0e-4 + weight_decay: 1.0e-2 + +# --- lr scheduler --- +lr_scheduler: + class: ConstantLR + kwargs: + factor: 1.0 + total_iters: 0 + +# --- training --- +train: + num_epochs: 50 + batch_size: 16 + shuffle: true + checkpoint_every_n_epochs: 5 + mixed_precision: fp16 + gradient_accumulation_steps: 1 + dataloader_num_workers: 4 + max_grad_norm: 1.0 + conditioning: discrete # 'discrete' (class_labels) or 'continuous' (encoder_hidden_states) + cfg_dropout: 0.0 # fraction of labels dropped for CFG training (0 disables) + ema_sigma_rels: [.02, .10] # e.g. [0.02, 0.10] to enable post-hoc EMA tracking + ema_update_every: 1 # optimizer steps between EMA updates (Karras assumes per-step) + ema_burn_in: 1000 # skip first N optimizer steps before starting EMA tracking + min_snr_gamma: null # set to e.g. 5.0 to enable Min-SNR loss weighting (Hang et al. 2023) + sigma_log_normal: null # set to e.g. [-1.2, 1.2] for EDM-style log-σ ~ Normal(P_mean, P_std) sampling + verbose: true + force_cpu: false + pin_memory: false + +# --- generation --- +generate: + scheduler: null # new scheduler for inference, default is train scheduler + num_steps: null # number of inference steps, default is training steps + s_churn: null # EDM stochasticity (Euler/Heun-family only); 0=ODE, larger=more SDE-like + s_tmin: null # restrict churn to t >= s_tmin + s_tmax: null # restrict churn to t <= s_tmax + s_noise: null # noise magnitude multiplier during churn (default 1.0) + n_samples: 64 # total number of samples to generate + batch_size: null # samples per forward pass; null → single batch of n_samples + image_shape: null # e.g. [1, 64, 64]; null → inferred from model.config + conditioning: discrete # 'discrete' or 'continuous' (must match training) + labels: null # discrete: list of int class labels (length 1 or n_samples) + continuous_labels: null # continuous: path to .npy of shape (n_samples, D) or (1, D) + guidance_scale: null # CFG guidance scale; null disables amplification + ema_sigma_rel: null # synthesize EMA at this target before sampling; null disables + seed: null # set for reproducibility + device: null # null → cuda if available else cpu diff --git a/cosmodiff/data/params_Illustris.txt b/cosmodiff/data/params_Illustris.txt new file mode 100644 index 0000000..80d682f --- /dev/null +++ b/cosmodiff/data/params_Illustris.txt @@ -0,0 +1,34 @@ +3.089999999999999969e-01 9.789999999999999813e-01 3.112340000000000106e+00 1.121939999999999937e+00 6.684999999999999831e-01 5.318199999999999594e-01 +1.310000000000000053e-01 7.126000000000000112e-01 1.321339999999999959e+00 3.195200000000000262e-01 1.701729999999999965e+00 5.073299999999999477e-01 +3.042000000000000259e-01 9.866000000000000325e-01 9.767099999999999671e-01 1.543279999999999985e+00 8.083200000000000385e-01 5.747499999999999831e-01 +2.358000000000000096e-01 7.842000000000000082e-01 1.189210000000000100e+00 2.500119999999999898e+00 1.249200000000000088e+00 9.599299999999999500e-01 +1.174000000000000044e-01 8.913999999999999702e-01 1.853180000000000049e+00 5.896799999999999820e-01 8.906899999999999817e-01 1.659790000000000099e+00 +3.038000000000000145e-01 7.149999999999999689e-01 1.185910000000000020e+00 4.897099999999999786e-01 6.403799999999999493e-01 1.376500000000000057e+00 +3.245999999999999996e-01 9.769999999999999796e-01 9.138300000000000312e-01 1.617760000000000087e+00 6.901599999999999957e-01 8.218800000000000550e-01 +2.713999999999999746e-01 9.337999999999999634e-01 3.629899999999999793e-01 7.463899999999999979e-01 7.732499999999999929e-01 1.687629999999999963e+00 +1.181999999999999995e-01 8.030000000000000471e-01 5.594200000000000284e-01 3.773599999999999732e-01 1.927859999999999907e+00 6.787699999999999845e-01 +4.390000000000000013e-01 7.897999999999999465e-01 6.232999999999999652e-01 3.168799999999999950e-01 1.469170000000000087e+00 5.983199999999999630e-01 +2.933999999999999941e-01 8.901999999999999913e-01 9.447499999999999787e-01 3.535709999999999908e+00 9.519800000000000484e-01 9.493399999999999617e-01 +1.242000000000000048e-01 9.409999999999999476e-01 9.317400000000000126e-01 5.441200000000000481e-01 1.007649999999999935e+00 6.466199999999999726e-01 +4.082000000000000073e-01 7.621999999999999886e-01 1.905280000000000085e+00 1.772759999999999891e+00 5.296100000000000252e-01 1.111879999999999979e+00 +4.981999999999999762e-01 9.649999999999999689e-01 2.679570000000000007e+00 2.675699999999999745e-01 1.380319999999999991e+00 5.684099999999999708e-01 +3.130000000000000004e-01 6.754000000000000004e-01 2.932100000000000262e-01 5.307100000000000151e-01 7.174700000000000522e-01 5.194299999999999473e-01 +1.650000000000000078e-01 9.761999999999999567e-01 4.924299999999999788e-01 1.176090000000000080e+00 6.220099999999999518e-01 1.298640000000000017e+00 +4.550000000000000155e-01 7.822000000000000064e-01 5.277800000000000269e-01 3.103720000000000034e+00 1.293250000000000011e+00 8.615500000000000380e-01 +1.865999999999999881e-01 6.538000000000000478e-01 1.041020000000000056e+00 1.714749999999999996e+00 5.108599999999999808e-01 5.252200000000000202e-01 +1.365999999999999992e-01 9.973999999999999533e-01 4.965499999999999914e-01 1.958839999999999915e+00 1.340710000000000068e+00 1.601029999999999953e+00 +2.066000000000000059e-01 6.790000000000000480e-01 2.070530000000000204e+00 1.159899999999999931e+00 1.578989999999999894e+00 1.121170000000000000e+00 +4.818000000000000060e-01 6.350000000000000089e-01 3.190980000000000150e+00 2.773900000000000254e-01 8.544100000000000028e-01 1.040300000000000002e+00 +1.701999999999999902e-01 6.985999999999999988e-01 3.186399999999999788e-01 3.073600000000000221e-01 1.175280000000000102e+00 7.510599999999999499e-01 +4.461999999999999855e-01 9.614000000000000323e-01 6.425999999999999490e-01 7.977499999999999591e-01 1.993079999999999963e+00 5.151299999999999768e-01 +1.781999999999999973e-01 7.137999999999999901e-01 6.624999999999999778e-01 2.977289999999999992e+00 7.285000000000000364e-01 6.475199999999999845e-01 +4.173999999999999932e-01 8.541999999999999593e-01 2.025110000000000188e+00 3.863750000000000018e+00 1.147899999999999920e+00 7.531499999999999861e-01 +3.794000000000000150e-01 8.045999999999999819e-01 2.785620000000000207e+00 1.295940000000000092e+00 9.773899999999999810e-01 1.033109999999999973e+00 +3.674000000000000044e-01 8.165999999999999925e-01 5.547800000000000509e-01 2.358709999999999862e+00 7.531499999999999861e-01 1.537940000000000085e+00 +1.965999999999999970e-01 6.186000000000000387e-01 4.789700000000000069e-01 4.776400000000000090e-01 1.998609999999999998e+00 1.113420000000000076e+00 +1.925999999999999934e-01 9.065999999999999615e-01 1.658639999999999892e+00 1.257009999999999961e+00 1.211670000000000025e+00 7.552400000000000224e-01 +2.086000000000000076e-01 9.741999999999999549e-01 2.385019999999999918e+00 8.156400000000000317e-01 9.841800000000000548e-01 9.506599999999999495e-01 +2.217999999999999972e-01 8.546000000000000263e-01 4.395200000000000218e-01 1.635799999999999921e+00 1.821339999999999959e+00 1.277209999999999956e+00 +2.142000000000000015e-01 6.373999999999999666e-01 2.225299999999999834e+00 4.227899999999999991e-01 1.594379999999999908e+00 5.629199999999999759e-01 +1.426000000000000045e-01 7.026000000000000023e-01 4.146599999999999731e-01 1.070289999999999964e+00 5.422400000000000553e-01 7.862200000000000299e-01 +1.406000000000000028e-01 9.574000000000000288e-01 6.515699999999999825e-01 4.856500000000000261e-01 6.475199999999999845e-01 5.325499999999999678e-01 diff --git a/cosmodiff/optim.py b/cosmodiff/optim.py index 9b44196..308e61c 100644 --- a/cosmodiff/optim.py +++ b/cosmodiff/optim.py @@ -1,4 +1,5 @@ import os +import inspect import pickle import yaml import numpy as np @@ -24,6 +25,163 @@ def _to_yaml_safe(obj): return obj +def _is_flow_matching(scheduler) -> bool: + """True if ``scheduler`` is a flow-matching variant. + + Detects ``FlowMatchEulerDiscreteScheduler``, ``FlowMatchHeunDiscreteScheduler``, + and any future ``FlowMatch*`` scheduler in ``diffusers``. The branching in + :func:`train` keys off this to use linear interpolation + velocity target + instead of the DDPM-style ``add_noise`` + ``prediction_type`` switch. + """ + return type(scheduler).__name__.startswith('FlowMatch') + + +def sample_timesteps( + scheduler, + batch_size: int, + device: str = 'cpu', + sigma_log_normal: Optional[tuple] = None, + log_sigma_schedule: Optional[torch.Tensor] = None, +): + """Sample per-batch timesteps for training. + + Dispatches on the scheduler family. For flow-matching schedulers, + samples an interpolation parameter ``sigma`` in ``[0, 1]`` (uniformly + by default, or via a logit-normal if ``sigma_log_normal`` is given, + matching the SD3-style training distribution). For diffusion + schedulers, samples integer timesteps in ``[0, T)`` — uniformly by + default, or via EDM-style log-σ sampling backed by + ``log_sigma_schedule``. + + Args: + scheduler: A ``diffusers`` scheduler. Flow-matching variants are + detected by class name (``FlowMatch*``); all other schedulers + are treated as DDPM-style. Used to read + ``scheduler.config.num_train_timesteps`` (denoted ``T``) and + to determine the timestep dtype convention. + batch_size (int): Number of timesteps to sample (one per batch item). + device: Target device for the returned tensors (e.g. ``'cpu'``, + ``'cuda'``, or a ``torch.device``). + sigma_log_normal (tuple of float, optional): ``(P_mean, P_std)``. + Under flow matching, samples are drawn as + ``sigmoid(P_mean + P_std * randn)`` (SD3-style logit-normal + ``t`` distribution). Under diffusion, samples are drawn as + ``log σ ~ Normal(P_mean, P_std)`` and mapped to the nearest + discrete timestep via ``log_sigma_schedule``. ``None`` + (default) uses uniform sampling. + log_sigma_schedule (torch.Tensor, optional): Precomputed + ``log σ`` values for the diffusion scheduler at each integer + timestep (shape ``(T,)``, monotonically increasing in ``t``). + Required only when ``sigma_log_normal`` is set and the + scheduler is *not* flow matching; ignored otherwise. + + Returns: + tuple: ``(timesteps, sigma)``. + + * ``timesteps`` (torch.Tensor of shape ``(batch_size,)``): + for diffusion, an integer ``long`` tensor in ``[0, T)``; + for flow matching, a float tensor in ``[0, T)`` (equal to + ``sigma * T``). + * ``sigma`` (torch.Tensor of shape ``(batch_size,)`` or + ``None``): the FM interpolant in ``[0, 1]`` when + ``scheduler`` is flow matching; ``None`` for diffusion + schedulers (which don't expose a sigma at this stage). + """ + T = scheduler.config.num_train_timesteps + if _is_flow_matching(scheduler): + # FM: sample sigma in [0, 1] (optionally logit-normal à la SD3) + if sigma_log_normal is not None: + P_mean, P_std = sigma_log_normal + sigma = torch.sigmoid(P_mean + P_std * torch.randn(batch_size, device=device)) + else: + sigma = torch.rand(batch_size, device=device) + timesteps = (sigma * T).clamp(0, T - 1) + return timesteps, sigma + + # diffusion branch + if sigma_log_normal is not None and log_sigma_schedule is not None: + P_mean, P_std = sigma_log_normal + log_sigma = P_mean + P_std * torch.randn(batch_size, device=device) + sched = log_sigma_schedule.to(device) + log_sigma = torch.clamp(log_sigma, sched.min().item(), sched.max().item()) + timesteps = torch.searchsorted(sched, log_sigma) + timesteps = torch.clamp(timesteps, 0, T - 1).long() + else: + timesteps = torch.randint(0, T, (batch_size,), device=device).long() + return timesteps, None + + +def noise_and_target( + scheduler, + samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.Tensor, + sigma: Optional[torch.Tensor] = None, +): + """Compute ``(noisy_samples, target)`` for one training step. + + Dispatches on the scheduler family. For flow-matching schedulers the + noisy sample is a linear interpolation between data and noise via + ``sigma``, and the target is the constant velocity field + ``noise - samples``. For diffusion schedulers the noisy sample is + produced by ``scheduler.add_noise`` and the target depends on + ``scheduler.config.prediction_type`` (``epsilon``, ``v_prediction``, + or ``sample``). + + Args: + scheduler: A ``diffusers`` scheduler. Flow-matching variants are + detected by class name (``FlowMatch*``). For diffusion + schedulers, ``scheduler.config.prediction_type`` selects the + target form and ``scheduler.add_noise`` / + ``scheduler.get_velocity`` provide the noised sample and + ``v_prediction`` target respectively. + samples (torch.Tensor): Clean training samples of shape + ``(B, ...)``. For 2D images this is typically + ``(B, C, H, W)``. + noise (torch.Tensor): Gaussian noise tensor with the same shape + and dtype as ``samples`` (typically drawn from + ``torch.randn_like(samples)``). + timesteps (torch.Tensor): Per-batch timesteps of shape ``(B,)`` + as returned by :func:`sample_timesteps`. Integer dtype for + diffusion, float for flow matching. + sigma (torch.Tensor, optional): The FM interpolant of shape + ``(B,)`` in ``[0, 1]``, as returned by :func:`sample_timesteps` + for flow-matching schedulers. Required when ``scheduler`` is + flow matching; ignored otherwise. + + Returns: + tuple: ``(noisy_samples, target)``, both ``torch.Tensor`` with the + same shape and dtype as ``samples``. + + * ``noisy_samples``: input fed to the model (the network's + ``x_t``). + * ``target``: the regression target for the MSE loss. + + Raises: + ValueError: If ``scheduler`` is a diffusion scheduler with an + unsupported ``prediction_type``. + """ + if _is_flow_matching(scheduler): + # Flow Matching + sigma_b = sigma.view(-1, *[1] * (samples.ndim - 1)) + noisy = sigma_b * noise + (1.0 - sigma_b) * samples + target = noise - samples + return noisy, target + + # DDPM + noisy = scheduler.add_noise(samples, noise, timesteps) + prediction_type = scheduler.config.prediction_type + if prediction_type == 'epsilon': + target = noise + elif prediction_type == 'v_prediction': + target = scheduler.get_velocity(samples, noise, timesteps) + elif prediction_type == 'sample': + target = samples + else: + raise ValueError(f"Unsupported prediction_type: {prediction_type}") + return noisy, target + + def train( dataset, model=None, @@ -43,6 +201,13 @@ def train( max_grad_norm: float = 1.0, force_cpu: bool = False, pin_memory: bool = False, + cfg_dropout: float = 0.0, + conditioning: str = 'discrete', + ema_sigma_rels: Optional[tuple] = None, + ema_update_every: int = 1, + ema_burn_in: int = 0, + min_snr_gamma: Optional[float] = None, + sigma_log_normal: Optional[tuple] = None, verbose: bool = True, ): """Train a diffusers diffusion model. @@ -98,6 +263,56 @@ def train( ``1.0``. force_cpu (bool): Force the model to the CPU even if CUDA available pin_memory (bool): pin dataset memory if on CPU + cfg_dropout (float): Fraction of training batches where labels are + replaced with a null token (``class_labels`` mode) or zeros + (``encoder_hidden_states`` mode) for classifier-free guidance + training. Set to ``0.0`` (default) to disable CFG. + conditioning (str): How labels are passed to the model. + ``'discrete'`` (default) passes integer labels via + ``class_labels=`` (e.g. ``UNet2DModel`` with ``num_class_embeds`` + or ``DiTTransformer2DModel``). ``'continuous'`` passes float + labels via ``encoder_hidden_states=`` (e.g. + ``UNet2DConditionModel``); labels of shape ``(B, D)`` are + automatically unsqueezed to ``(B, 1, D)``. + ema_sigma_rels (tuple of float, optional): If set, enables post-hoc + EMA tracking via ``ema-pytorch``. Two values are required (e.g. + ``(0.03, 0.15)``); they control the two power-function EMA + profiles checkpointed at each epoch. Choose them to bracket the + target ``sigma_rel`` range you want to synthesize, with the lower + value anchoring your floor and the upper value giving some + headroom above your ceiling. After training, + ``ema.synthesize_ema_model(sigma_rel=...)`` reconstructs any + target profile from those snapshots. ``None`` (default) disables + EMA entirely. + ema_update_every (int): How often (in optimizer steps) to update the + EMA profiles. Defaults to ``1`` (every step), matching the Karras + et al. 2024 reference implementation and the broader diffusion + training literature. Larger values silently distort the EMA + profile because the time-varying beta formula assumes per-step + updates — only increase if profiling shows the EMA lerp is a + meaningful fraction of step time. + ema_burn_in (int): Number of initial optimizer steps during which the + EMA is *not* updated. Defaults to ``0`` (start tracking from step + 0). When set, the EMA's internal time origin shifts to step + ``ema_burn_in``, so ``sigma_rel`` is interpreted as a fraction of + the post-burn-in trajectory ``[ema_burn_in, T_total]``. Useful to + avoid contaminating the EMA with the very-early random/unstable + weights, particularly for wider EMA profiles. + min_snr_gamma (float, optional): If set, applies Min-SNR loss + weighting from Hang et al. 2023. Recommended value is ``5.0``. + The per-sample loss is multiplied by ``min(SNR(t), gamma)`` + divided by a parameterization-dependent factor (``SNR(t)`` for + epsilon-prediction, ``SNR(t)+1`` for v-prediction, ``1`` for + sample-prediction), which balances training signal across + timesteps and is especially helpful for v-prediction. ``None`` + (default) disables Min-SNR (uniform MSE). + sigma_log_normal (tuple of float, optional): If set to a + ``(P_mean, P_std)`` tuple, samples log-σ from ``Normal(P_mean, + P_std)`` each step (Karras et al. EDM 2022) instead of uniform + timestep sampling. Sampled log-σ is mapped to the nearest + discrete timestep of the underlying scheduler via + ``torch.searchsorted``. EDM defaults are ``(-1.2, 1.2)``. + ``None`` (default) keeps standard uniform timestep sampling. verbose (bool): Print training progress and checkpoint messages. Defaults to ``True``. @@ -207,6 +422,33 @@ def _load_model_hook(models, input_dir): if resume_from_checkpoint is not None: accelerator.load_state(resume_from_checkpoint) + # post-hoc EMA (fresh each run; not restored on resume) + ema = None + if ema_sigma_rels is not None: + from pathlib import Path + from ema_pytorch import PostHocEMA + ema = PostHocEMA( + accelerator.unwrap_model(model), + sigma_rels=ema_sigma_rels, + checkpoint_every_num_steps='manual', + checkpoint_folder=Path(output_dir) / 'ema', + update_every=ema_update_every, + ) + + # null token for CFG label dropout (class_labels mode only) + cfg_null_token = None + if cfg_dropout > 0.0 and conditioning == 'discrete': + cfg_null_token = accelerator.unwrap_model(model).config.num_class_embeds - 1 + + # precompute log-σ schedule for EDM-style log-normal sigma sampling + # (only meaningful for diffusion schedulers; FM uses sigma_log_normal as + # logit-normal `t` parameters directly — no schedule lookup needed) + log_sigma_schedule = None + if sigma_log_normal is not None and not _is_flow_matching(noise_scheduler): + alphas_cumprod = noise_scheduler.alphas_cumprod.clamp(min=1e-10) + log_sigma_schedule = 0.5 * torch.log((1.0 - alphas_cumprod) / alphas_cumprod) + # monotonically increasing in t; we'll map sampled log-σ to nearest t via searchsorted + # ------------------------------------------------------------------ # # 5. Training loop # # ------------------------------------------------------------------ # @@ -235,37 +477,64 @@ def _load_model_hook(models, input_dir): for batch in progress: images = batch["images"] labels = batch.get("labels", None) + if labels is not None and cfg_dropout > 0.0: + drop_mask = torch.rand(labels.shape[0], device=labels.device) < cfg_dropout + if conditioning == 'discrete': + labels = torch.where(drop_mask, torch.full_like(labels, cfg_null_token), labels) + else: + mask = drop_mask.view(-1, *([1] * (labels.ndim - 1))) + labels = torch.where(mask, torch.zeros_like(labels), labels) t0 = time.time() with accelerator.accumulate(model): - # --- forward diffusion ---------------------------------- + # --- forward diffusion (or FM linear interpolation) ----- noise = torch.randn_like(images) - timesteps = torch.randint( - 0, - noise_scheduler.config.num_train_timesteps, - (images.shape[0],), + timesteps, sigma = sample_timesteps( + noise_scheduler, + images.shape[0], device=images.device, - ).long() - noisy_images = noise_scheduler.add_noise(images, noise, timesteps) + sigma_log_normal=sigma_log_normal, + log_sigma_schedule=log_sigma_schedule, + ) + noisy_images, target = noise_and_target( + noise_scheduler, images, noise, timesteps, sigma=sigma, + ) # --- predict noise -------------------------------------- with accelerator.autocast(): if labels is not None: - noise_pred = model( - noisy_images, - timestep=timesteps, - class_labels=labels, - return_dict=False, - )[0] + if conditioning == 'discrete': + cond_kwargs = {'class_labels': labels} + else: + cond = labels.float() + if cond.ndim == 2: + cond = cond.unsqueeze(1) + cond_kwargs = {'encoder_hidden_states': cond} + pred = model(noisy_images, timestep=timesteps, return_dict=False, **cond_kwargs)[0] else: - noise_pred = model( - noisy_images, - timesteps, - return_dict=False, - )[0] + pred = model(noisy_images, timesteps, return_dict=False)[0] - loss = F.mse_loss(noise_pred, noise) + if min_snr_gamma is None or _is_flow_matching(noise_scheduler): + # Min-SNR is a DDPM-style weighting and does not apply + # to flow matching — use plain MSE in that case. + loss = F.mse_loss(pred, target) + else: + # Min-SNR loss weighting (Hang et al. 2023). SNR = α_bar / (1 - α_bar). + prediction_type = noise_scheduler.config.prediction_type + alphas_cumprod = noise_scheduler.alphas_cumprod.to(timesteps.device) + snr = alphas_cumprod[timesteps] / (1.0 - alphas_cumprod[timesteps]) + clipped_snr = torch.clamp(snr, max=min_snr_gamma) + if prediction_type == 'epsilon': + weight = clipped_snr / snr + elif prediction_type == 'sample': + weight = clipped_snr + else: # v_prediction + weight = clipped_snr / (snr + 1.0) + per_sample_mse = F.mse_loss(pred, target, reduction='none').mean( + dim=list(range(1, pred.ndim)) + ) + loss = (weight * per_sample_mse).mean() # --- backward ------------------------------------------- accelerator.backward(loss) @@ -275,6 +544,9 @@ def _load_model_hook(models, input_dir): optimizer.step() optimizer.zero_grad(set_to_none=True) + if ema is not None and global_step >= ema_burn_in: + ema.update() + lr_scheduler.step() batch_time = time.time() - t0 batch_loss = loss.detach().item() @@ -327,6 +599,8 @@ def _load_model_hook(models, input_dir): "class": f"{raw_sched.__class__.__module__}.{raw_sched.__class__.__name__}", "kwargs": utils._get_lr_scheduler_kwargs(raw_sched), }, + "ema_sigma_rels": list(ema_sigma_rels) if ema_sigma_rels is not None else None, + "ema_burn_in": ema_burn_in, } with open(os.path.join(ckpt_save_path, "checkpoint_config.yaml"), "w") as f: yaml.dump(_to_yaml_safe(ckpt_cfg), f) @@ -334,6 +608,12 @@ def _load_model_hook(models, input_dir): # Model weights (via hook) + optimizer moments + grad scaler + RNG accelerator.save_state(ckpt_save_path) + if ema is not None and ema.step.item() > 0: + from pathlib import Path + ema.checkpoint_folder = Path(ckpt_save_path) / 'ema' + ema.checkpoint_folder.mkdir(exist_ok=True) + ema.checkpoint() + if hasattr(dataset, "augmentations") and dataset.augmentations is not None: with open(os.path.join(ckpt_save_path, "augmentations.pkl"), "wb") as f: pickle.dump(dataset.augmentations, f) @@ -345,7 +625,7 @@ def _load_model_hook(models, input_dir): accelerator.print(f"Checkpoint saved → {ckpt_save_path}") accelerator.end_training() - return metrics + return {'metrics': metrics, 'ema': ema} @torch.no_grad() @@ -355,16 +635,23 @@ def generate( batch_size: int = 1, image_shape: tuple[int, ...] = (1, 64, 64), labels: Optional[torch.Tensor] = None, - ddim_thinning: Optional[int] = None, + guidance_scale: Optional[float] = None, + conditioning: str = 'discrete', + num_steps: Optional[int] = None, + s_churn: Optional[float] = None, + s_tmin: Optional[float] = None, + s_tmax: Optional[float] = None, + s_noise: Optional[float] = None, renorm: Optional[Callable] = None, device: Optional[torch.device] = None, generator: Optional[torch.Generator] = None, ) -> torch.Tensor: """Generate a batch of novel images via reverse diffusion. - Compatible with DDPMScheduler and DDIMScheduler. The number of inference - steps defaults to ``noise_scheduler.config.num_train_timesteps``, or can - be reduced via ``ddim_thinning`` for faster DDIM sampling. + Compatible with any ``diffusers`` scheduler exposing + ``set_timesteps`` / ``step``. The number of inference steps defaults to + ``noise_scheduler.config.num_train_timesteps`` and can be reduced via + ``num_steps`` for faster sampling (DDIM, DPM-Solver, Heun, etc.). Args: model: Trained diffusers model (UNet2DModel or DiTTransformer2DModel). @@ -372,13 +659,39 @@ def generate( any compatible scheduler with ``set_timesteps`` and ``step``). batch_size (int): Number of images to generate. image_shape (tuple): Shape of each image ``(C, H, W)``. - labels (array-like, optional): Class indices of length ``batch_size``. - Required for class-conditional models (e.g. DiTTransformer2DModel); - ignored otherwise. - ddim_thinning (int, optional): Thinning factor relative to - ``num_train_timesteps``. E.g. ``ddim_thinning=10`` with 1000 - training steps runs 100 inference steps. Only meaningful with - DDIMScheduler; ignored (``None``) for DDPM. + labels (array-like, optional): Conditioning labels of length + ``batch_size``. Integer class indices for ``'class_labels'`` mode; + float array of shape ``(batch_size, D)`` for + ``'encoder_hidden_states'`` mode. + guidance_scale (float, optional): Classifier-free guidance scale. + ``None`` (default) uses the plain conditional prediction with no + amplification. A float in ``[1.0, 15.0]`` runs a double forward + pass with null labels and controls class adherence: + ``uncond + scale * (cond - uncond)``. + conditioning (str): Matches the mode used during training. + ``'discrete'`` (default) for integer class labels; + ``'continuous'`` for continuous float labels. + num_steps (int, optional): Number of inference steps to run. Passed + directly to ``noise_scheduler.set_timesteps``. Defaults to + ``noise_scheduler.config.num_train_timesteps`` (full schedule). + Use a smaller value with a higher-order solver (DDIM, DPM-Solver, + Heun, etc.) for fast sampling. + s_churn (float, optional): EDM-style stochasticity injection + (Karras et al. 2022). ``0`` (the scheduler default) gives pure + ODE sampling; larger values inject more noise at each step, + interpolating toward SDE-like behavior. Only consumed by + schedulers that accept it (Euler/Heun-family, including the + flow-matching variants); silently dropped for schedulers that + don't (DDPM, DDIM, DPM-Solver, etc.). + s_tmin (float, optional): Lower-bound timestep for churn gating. + Churn is only applied when ``t >= s_tmin``. Same compatibility + rules as ``s_churn``. + s_tmax (float, optional): Upper-bound timestep for churn gating. + Churn is only applied when ``t <= s_tmax``. Same compatibility + rules as ``s_churn``. + s_noise (float, optional): Multiplier on the magnitude of injected + noise during churn (default ``1.0`` in diffusers). Same + compatibility rules as ``s_churn``. renorm (callable, optional): Applied to the output tensor to convert images back to their original range (e.g. inverse of the normalization used at training time). @@ -396,7 +709,7 @@ def generate( model.eval() num_train_timesteps = noise_scheduler.config.num_train_timesteps - num_inference_steps = num_train_timesteps // ddim_thinning if ddim_thinning is not None else num_train_timesteps + num_inference_steps = num_steps if num_steps is not None else num_train_timesteps noise_scheduler.set_timesteps(num_inference_steps) images = torch.randn( @@ -406,22 +719,62 @@ def generate( ) if labels is not None: - labels = torch.as_tensor(labels, dtype=torch.long, device=device) + if conditioning == 'discrete': + labels = torch.as_tensor(labels, dtype=torch.long, device=device) + else: + labels = torch.as_tensor(labels, dtype=torch.float, device=device) + if labels.ndim == 2: + labels = labels.unsqueeze(1) # (B, D) -> (B, 1, D) + + null_labels = None + if labels is not None and guidance_scale is not None: + if conditioning == 'discrete': + null_token = model.config.num_class_embeds - 1 + null_labels = torch.full_like(labels, null_token) + else: + null_labels = torch.zeros_like(labels) + + # Filter EDM-style stochasticity kwargs to those this scheduler's + # `step()` actually accepts (Euler/Heun-family, including FM variants). + # Unsupported keys are silently dropped so the same generate() call + # works across all scheduler types. + step_supported = set(inspect.signature(noise_scheduler.step).parameters) + step_kwargs = { + k: v for k, v in ( + ('s_churn', s_churn), ('s_tmin', s_tmin), + ('s_tmax', s_tmax), ('s_noise', s_noise), + ) + if v is not None and k in step_supported + } for t in tqdm(noise_scheduler.timesteps, desc="Sampling"): - timesteps = torch.full((batch_size,), t, device=device, dtype=torch.long) + # FM schedulers use float timesteps; DDPM-family use int — preserve dtype. + timesteps = torch.full((batch_size,), t.item() if hasattr(t, 'item') else t, + device=device, dtype=t.dtype if hasattr(t, 'dtype') else torch.long) + + # Rescale the state into the form the model expects (identity for + # VP/DDIM/DPM-Solver/FlowMatch, division by √(σ²+1) for Euler/Heun + # on VP-trained models). Pass scaled to the model; pass unscaled + # `images` to scheduler.step which integrates in its native space. + # FlowMatch schedulers don't define scale_model_input at all, so + # we fall through to identity in that case. + if hasattr(noise_scheduler, 'scale_model_input'): + images_input = noise_scheduler.scale_model_input(images, t) + else: + images_input = images if labels is not None: - noise_pred = model( - images, - timestep=timesteps, - class_labels=labels, - return_dict=False, - )[0] + cond_key = 'class_labels' if conditioning == 'discrete' else 'encoder_hidden_states' + pred = model(images_input, timestep=timesteps, return_dict=False, **{cond_key: labels})[0] + if null_labels is not None: + uncond_pred = model(images_input, timestep=timesteps, return_dict=False, **{cond_key: null_labels})[0] + pred = uncond_pred + guidance_scale * (pred - uncond_pred) else: - noise_pred = model(images, timesteps, return_dict=False)[0] + pred = model(images_input, timesteps, return_dict=False)[0] - images = noise_scheduler.step(noise_pred, t, images, generator=generator).prev_sample + images = noise_scheduler.step( + pred, t, images, generator=generator, **step_kwargs, + ).prev_sample if renorm is not None: images = renorm(images) @@ -429,6 +782,204 @@ def generate( return images +def gamma_from_sigma_rel(sigma_rel: float) -> float: + """Convert a sigma_rel value to the corresponding power-function EMA exponent gamma. + + The power-function EMA assigns weight ``w(x) = (gamma+1) * x^gamma`` to + training progress ``x = s/t`` in ``[0, 1]``. ``sigma_rel`` is the + standard deviation of that distribution and is also equal to the relative + EMA length (fraction of training): ``sigma_rel=0.1`` means the EMA draws + its weight from roughly the last 10% of training (Karras et al. 2024, + Sec. 3.1). + + Args: + sigma_rel: relative standard deviation of the EMA kernel (e.g. 0.05 or 0.28). + + Returns: + Corresponding gamma exponent. + """ + from scipy.optimize import brentq + + def _f(gamma): + mean = (gamma + 1) / (gamma + 2) + second_moment = (gamma + 1) / (gamma + 3) + return np.sqrt(second_moment - mean ** 2) - sigma_rel + + return brentq(_f, 1e-4, 1e4) + + +def compute_ema_profiles( + sigma_rels_train: tuple, + checkpoint_epochs: list, + total_epochs: int, + sigma_rel_target: float, +) -> dict: + """Compute per-checkpoint and synthesized EMA weight profiles. + + Implements the analysis behind Fig. 4 of Karras et al. 2024. For each + ``(checkpoint_epoch, sigma_rel)`` pair a basis profile is computed; the + synthesized profile is the non-negative least-squares combination that + best approximates the target profile. + + Args: + sigma_rels_train: sigma_rel values used during training (e.g. ``(0.05, 0.28)``). + checkpoint_epochs: epochs at which ``ema.checkpoint()`` was called. + total_epochs: total training length (defines the target profile). + sigma_rel_target: target sigma_rel to synthesize (e.g. ``0.15``). + + Returns: + dict with keys: + ``'epochs'``: ``np.ndarray`` of shape ``(total_epochs,)`` + ``'basis'``: list of ``(checkpoint_epoch, sigma_rel, weights)`` tuples + ``'target'``: ``np.ndarray`` of shape ``(total_epochs,)`` + ``'synthesized'``: ``np.ndarray`` of shape ``(total_epochs,)`` + ``'coefficients'``: ``np.ndarray`` of shape ``(n_basis,)`` + """ + from scipy.optimize import nnls + + epochs = np.arange(1, total_epochs + 1) + gammas_train = [gamma_from_sigma_rel(s) for s in sigma_rels_train] + gamma_target = gamma_from_sigma_rel(sigma_rel_target) + + basis = [] + for t_i in checkpoint_epochs: + for gamma_j, s_j in zip(gammas_train, sigma_rels_train): + w = np.where(epochs <= t_i, (epochs / t_i) ** gamma_j, 0.0) + if w.sum() > 0: + w = w / w.sum() + basis.append((t_i, s_j, w)) + + w_target = (epochs / total_epochs) ** gamma_target + w_target = w_target / w_target.sum() + + B = np.column_stack([w for _, _, w in basis]) + coeffs, _ = nnls(B, w_target) + if coeffs.sum() > 0: + coeffs = coeffs / coeffs.sum() + w_synth = B @ coeffs + + return { + 'epochs': epochs, + 'basis': basis, + 'target': w_target, + 'synthesized': w_synth, + 'coefficients': coeffs, + } + + +def synthesize_ema_from_checkpoints( + model: torch.nn.Module, + output_dir: str, + sigma_rel_target: float, + sigma_rels: Optional[tuple] = None, + up_to_epoch: Optional[int] = None, +) -> torch.nn.Module: + """Synthesize a post-hoc EMA model by pooling snapshots across checkpoints. + + Gathers all ``.pt`` EMA snapshot files from every ``checkpoint-epoch-*/ema/`` + subdirectory in ``output_dir`` into a single temporary directory (via + symlinks) and delegates synthesis to ``PostHocEMA``. + + ``sigma_rels`` is read automatically from ``checkpoint_config.yaml`` if not + supplied explicitly. + + Args: + model: the nn.Module used during training (needed for parameter structure). + output_dir: root training directory containing ``checkpoint-epoch-*`` dirs. + sigma_rel_target: target EMA profile to synthesize, e.g. ``0.15``. + sigma_rels: sigma_rel values used during training, e.g. ``(0.05, 0.28)``. + If ``None``, read from the latest checkpoint's ``checkpoint_config.yaml``. + up_to_epoch: if set, only use checkpoints at or before this epoch number. + + Returns: + Synthesized nn.Module with EMA weights. + """ + import tempfile + from pathlib import Path + from ema_pytorch import PostHocEMA + + ckpt_dirs = sorted(Path(output_dir).glob('checkpoint-epoch-*')) + if up_to_epoch is not None: + ckpt_dirs = [d for d in ckpt_dirs if int(d.name.split('-')[-1]) <= up_to_epoch] + + if not ckpt_dirs: + raise ValueError(f"No checkpoints found in {output_dir}") + + if sigma_rels is None: + cfg_path = ckpt_dirs[-1] / 'checkpoint_config.yaml' + with open(cfg_path) as f: + ckpt_cfg = yaml.safe_load(f) + sigma_rels = ckpt_cfg.get('ema_sigma_rels') + if sigma_rels is None: + raise ValueError( + "sigma_rels not found in checkpoint_config.yaml — " + "pass sigma_rels explicitly or retrain with ema_sigma_rels set." + ) + sigma_rels = tuple(sigma_rels) + + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + for ckpt_dir in ckpt_dirs: + for pt_file in (ckpt_dir / 'ema').glob('*.pt'): + (tmp_path / pt_file.name).symlink_to(pt_file.resolve()) + + ema = PostHocEMA( + model, + sigma_rels=sigma_rels, + checkpoint_folder=tmp_path, + checkpoint_every_num_steps='manual', + ) + return ema.synthesize_ema_model(sigma_rel=sigma_rel_target) + + +def load_ema_snapshot( + model: torch.nn.Module, + checkpoint_dir: str, + profile_index: int = 0, +) -> torch.nn.Module: + """Load a single raw EMA snapshot into ``model`` in-place. + + Reads the ``.pt`` snapshot file for the given training profile from + ``/ema/`` and copies its weights into ``model`` (no + synthesis, no pooling across checkpoints). + + Use this when you want exactly the EMA at one of the trained ``sigma_rels`` + captured at one specific epoch. For arbitrary target ``sigma_rel`` values + or for pooling across many checkpoints, use + :func:`synthesize_ema_from_checkpoints` instead. + + Args: + model: the nn.Module to load weights into; modified in place and returned. + checkpoint_dir: a ``checkpoint-epoch-XXXX`` directory containing an + ``ema/`` subdirectory. + profile_index: which trained profile to load — ``0`` for + ``sigma_rels[0]`` (e.g. 0.05), ``1`` for ``sigma_rels[1]`` + (e.g. 0.28). Defaults to ``0``. + + Returns: + The same ``model``, with EMA weights loaded. + """ + from pathlib import Path + + ema_dir = Path(checkpoint_dir) / 'ema' + pt_files = sorted(ema_dir.glob(f'{profile_index}.*.pt')) + if not pt_files: + raise ValueError( + f"No EMA snapshots for profile {profile_index} in {ema_dir}" + ) + # filename format is {profile_index}.{global_step}.pt — pick the latest step + pt_file = max(pt_files, key=lambda p: int(p.stem.split('.')[1])) + + snap = torch.load(pt_file, map_location='cpu', weights_only=False) + ema_state = { + k[len('ema_model.'):]: v + for k, v in snap.items() + if k.startswith('ema_model.') + } + model.load_state_dict(ema_state) + return model + + class PCAEncoder(torch.nn.Module): """Linear PCA encoder as an nn.Module. diff --git a/cosmodiff/tests/test_optim.py b/cosmodiff/tests/test_optim.py index 8cf7fad..ef3db3b 100644 --- a/cosmodiff/tests/test_optim.py +++ b/cosmodiff/tests/test_optim.py @@ -2,10 +2,25 @@ import tempfile import numpy as np import torch -from diffusers import UNet2DModel, DDPMScheduler, DDIMScheduler, DiTTransformer2DModel +from diffusers import ( + UNet2DModel, + UNet2DConditionModel, + DDPMScheduler, + DDIMScheduler, + EulerDiscreteScheduler, + DiTTransformer2DModel, + PixArtTransformer2DModel, + FlowMatchEulerDiscreteScheduler, +) from cosmodiff.utils import load_checkpoint, ArrayDataset, find_latest_checkpoint -from cosmodiff.optim import train, generate, compute_fid, compute_kid, build_pca_encoder +from cosmodiff.optim import train, generate, compute_fid, compute_kid, build_pca_encoder, synthesize_ema_from_checkpoints, compute_ema_profiles, load_ema_snapshot from cosmodiff.augment import RandomRoll, RandomFlip +from cosmodiff.data import DATA_PATH + +SIM_PATH = DATA_PATH / 'IllustrisTNG_Mcdm.npy' +PARAMS_PATH = DATA_PATH / 'params_Illustris.txt' +# shape: (34,) simulations × 6 cosmological parameters (a,b,c,d,e,f) +N_SIMS, N_PARAMS = 34, 6 def test_train_basic(): @@ -22,12 +37,12 @@ def test_train_basic(): ) images = torch.randn(16, 1, 8, 8) - augmentations = torch.nn.Sequential(RandomRoll(dims=(-1,-2)), RandomFlip(dims=(-1, -2))) + augmentations = torch.nn.Sequential(RandomRoll(size=8, dims=(-1,-2)), RandomFlip(dims=(-1, -2))) dataset = ArrayDataset(images, augmentations=augmentations) with tempfile.TemporaryDirectory() as tmp_dir: initial_weights = model.conv_in.weight.data.clone() - metrics = train( + result = train( dataset, model, noise_scheduler=DDPMScheduler(num_train_timesteps=10), @@ -49,8 +64,8 @@ def test_train_basic(): assert os.path.exists(os.path.join(ckpt_path, "metrics.json")) # training checks: finite output, and weights changed - assert all(torch.isfinite(torch.tensor(v)) for v in metrics["loss"]) - assert all(torch.isfinite(torch.tensor(v)) for v in metrics["epoch_loss"]) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["loss"]) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["epoch_loss"]) assert not torch.allclose(model.conv_in.weight.data, initial_weights) # load_checkpoint @@ -66,7 +81,7 @@ def test_train_basic(): # continue training from checkpoint initial_weights = model.conv_in.weight.data.clone() - metrics = train( + result = train( dataset, resume_from_checkpoint=ckpt_path, num_epochs=2, @@ -86,8 +101,8 @@ def test_train_basic(): ) # training checks: finite output, and weights changed - assert all(torch.isfinite(torch.tensor(v)) for v in metrics["loss"]) - assert all(torch.isfinite(torch.tensor(v)) for v in metrics["epoch_loss"]) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["loss"]) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["epoch_loss"]) assert not torch.allclose(_model2.conv_in.weight.data, initial_weights) @@ -107,12 +122,12 @@ def test_train_conditional_dit(): images = torch.randn(16, 1, 4, 4) labels = torch.randint(0, 4, (16,)) - augmentations = torch.nn.Sequential(RandomRoll(dims=(-1,-2)), RandomFlip(dims=(-1, -2))) + augmentations = torch.nn.Sequential(RandomRoll(size=4, dims=(-1,-2)), RandomFlip(dims=(-1, -2))) dataset = ArrayDataset(images, labels=labels, augmentations=augmentations) with tempfile.TemporaryDirectory() as tmp_dir: initial_weights = model.transformer_blocks[0].ff.net[0].proj.weight.clone() - metrics = train( + result = train( dataset, model, noise_scheduler=DDPMScheduler(num_train_timesteps=10), @@ -134,8 +149,8 @@ def test_train_conditional_dit(): assert os.path.exists(os.path.join(ckpt_path, "metrics.json")) # training checks: finite output, and weights changed - assert all(torch.isfinite(torch.tensor(v)) for v in metrics["loss"]) - assert all(torch.isfinite(torch.tensor(v)) for v in metrics["epoch_loss"]) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["loss"]) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["epoch_loss"]) assert not torch.allclose(model.transformer_blocks[0].ff.net[0].proj.weight, initial_weights) # load_checkpoint @@ -163,6 +178,52 @@ def _make_unet(sample_size=8): ) +def _make_unet_class_cond(n_classes=4, sample_size=8): + return UNet2DModel( + sample_size=sample_size, + in_channels=1, + out_channels=1, + layers_per_block=1, + block_out_channels=(16, 16), + down_block_types=("DownBlock2D", "DownBlock2D"), + up_block_types=("UpBlock2D", "UpBlock2D"), + norm_num_groups=8, + num_class_embeds=n_classes + 1, # last index is the null token + ) + + +def _make_unet_condition(sample_size=8, n_params=2): + return UNet2DConditionModel( + sample_size=sample_size, + in_channels=1, + out_channels=1, + layers_per_block=1, + block_out_channels=(16, 16), + down_block_types=("DownBlock2D", "DownBlock2D"), + up_block_types=("UpBlock2D", "UpBlock2D"), + norm_num_groups=8, + cross_attention_dim=16, + encoder_hid_dim=n_params, + encoder_hid_dim_type="text_proj", + ) + + +def _make_pixart(sample_size=8, n_params=N_PARAMS): + return PixArtTransformer2DModel( + sample_size=sample_size, + patch_size=2, + in_channels=1, + out_channels=1, + num_layers=1, + num_attention_heads=2, + attention_head_dim=8, + cross_attention_dim=16, + caption_channels=n_params, + use_additional_conditions=False, + norm_num_groups=None, + ) + + def test_generate_ddpm(): """generate() returns the right shape and finite values with DDPMScheduler.""" model = _make_unet() @@ -173,11 +234,11 @@ def test_generate_ddpm(): assert torch.isfinite(images).all() -def test_generate_ddim_thinning(): - """ddim_thinning reduces inference steps and output is still valid.""" +def test_generate_num_steps(): + """num_steps reduces inference steps and output is still valid.""" model = _make_unet() scheduler = DDIMScheduler(num_train_timesteps=10) - images = generate(model, scheduler, batch_size=2, image_shape=(1, 8, 8), ddim_thinning=2) + images = generate(model, scheduler, batch_size=2, image_shape=(1, 8, 8), num_steps=5) assert images.shape == (2, 1, 8, 8) assert torch.isfinite(images).all() @@ -302,4 +363,710 @@ def test_kid_smaller_same_dist(): assert kid_same < kid_diff +# ------------------------------------------------------------------ # +# CFG / conditioning tests # +# ------------------------------------------------------------------ # + +def test_train_cfg_class_labels(): + """train() with cfg_dropout runs without error and produces finite loss.""" + n_classes = 4 + model = _make_unet_class_cond(n_classes=n_classes) + images = torch.randn(8, 1, 8, 8) + labels = torch.randint(0, n_classes, (8,)) + dataset = ArrayDataset(images, labels=labels) + + with tempfile.TemporaryDirectory() as tmp_dir: + result = train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=1, + batch_size=4, + checkpoint_every_n_epochs=2, + mixed_precision="no", + output_dir=tmp_dir, + force_cpu=True, + verbose=False, + cfg_dropout=0.5, + conditioning='discrete', + ) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["loss"]) + + +def test_train_encoder_hidden_states(): + """train() with conditioning='continuous' using real cosmological params.""" + params = np.loadtxt(PARAMS_PATH, dtype=np.float32) # (34, 6) + images = torch.randn(len(params), 1, 8, 8) + labels = torch.as_tensor(params) # (34, 6) + dataset = ArrayDataset(images, labels=labels) + model = _make_unet_condition(n_params=N_PARAMS) + + with tempfile.TemporaryDirectory() as tmp_dir: + result = train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=1, + batch_size=4, + checkpoint_every_n_epochs=2, + mixed_precision="no", + output_dir=tmp_dir, + force_cpu=True, + verbose=False, + conditioning='continuous', + ) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["loss"]) + + +def test_generate_encoder_hidden_states(): + """generate() with conditioning='continuous' using real cosmological params.""" + params = np.loadtxt(PARAMS_PATH, dtype=np.float32) # (34, 6) + model = _make_unet_condition(n_params=N_PARAMS) + scheduler = DDPMScheduler(num_train_timesteps=10) + # use first 3 simulations' parameters as conditioning + labels = torch.as_tensor(params[:3]) # (3, 6) + + images = generate( + model, scheduler, + batch_size=3, image_shape=(1, 8, 8), + labels=labels, + conditioning='continuous', + ) + assert images.shape == (3, 1, 8, 8) + assert torch.isfinite(images).all() + + +def test_generate_cfg_guidance_scale(): + """guidance_scale != 1.0 produces different output than guidance_scale=1.0.""" + n_classes = 4 + model = _make_unet_class_cond(n_classes=n_classes) + scheduler = DDPMScheduler(num_train_timesteps=10) + labels = torch.tensor([0, 1, 2]) + + g1 = torch.Generator().manual_seed(0) + out_no_cfg = generate( + model, scheduler, + batch_size=3, image_shape=(1, 8, 8), + labels=labels, guidance_scale=None, + conditioning='discrete', + generator=g1, + ) + + g2 = torch.Generator().manual_seed(0) + out_cfg = generate( + model, scheduler, + batch_size=3, image_shape=(1, 8, 8), + labels=labels, guidance_scale=2.0, + conditioning='discrete', + generator=g2, + ) + + assert out_no_cfg.shape == out_cfg.shape == (3, 1, 8, 8) + assert torch.isfinite(out_cfg).all() + assert not torch.allclose(out_no_cfg, out_cfg) + + +def test_train_pixart(): + """train() with PixArtTransformer2DModel and real cosmological params.""" + params = np.loadtxt(PARAMS_PATH, dtype=np.float32) # (34, 6) + images = torch.randn(len(params), 1, 8, 8) + labels = torch.as_tensor(params) + dataset = ArrayDataset(images, labels=labels) + model = _make_pixart() + + with tempfile.TemporaryDirectory() as tmp_dir: + result = train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=1, + batch_size=4, + checkpoint_every_n_epochs=2, + mixed_precision="no", + output_dir=tmp_dir, + force_cpu=True, + verbose=False, + conditioning='continuous', + ) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["loss"]) + + +def test_generate_pixart(): + """generate() with PixArtTransformer2DModel and real cosmological params.""" + params = np.loadtxt(PARAMS_PATH, dtype=np.float32) + model = _make_pixart() + scheduler = DDPMScheduler(num_train_timesteps=10) + labels = torch.as_tensor(params[:3]) # (3, 6) + + images = generate( + model, scheduler, + batch_size=3, image_shape=(1, 8, 8), + labels=labels, + conditioning='continuous', + ) + assert images.shape == (3, 1, 8, 8) + assert torch.isfinite(images).all() + + +# ------------------------------------------------------------------ # +# EMA tests # +# ------------------------------------------------------------------ # + +def test_train_ema(): + """train() with ema_sigma_rels checkpoints EMA profiles and returns ema object.""" + model = _make_unet() + images = torch.randn(8, 1, 8, 8) + dataset = ArrayDataset(images) + + with tempfile.TemporaryDirectory() as tmp_dir: + result = train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=2, + batch_size=4, + checkpoint_every_n_epochs=2, + mixed_precision="no", + output_dir=tmp_dir, + force_cpu=True, + verbose=False, + ema_sigma_rels=(0.05, 0.28), + ) + + assert result['ema'] is not None + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']["loss"]) + + # EMA profiles should be checkpointed inside the epoch checkpoint dir + ckpt_path = os.path.join(tmp_dir, "checkpoint-epoch-0001") + assert os.path.isdir(os.path.join(ckpt_path, 'ema')) + + # should be able to synthesize a new EMA profile post-hoc + synthesized = result['ema'].synthesize_ema_model(sigma_rel=0.15) + assert synthesized is not None + + +def test_train_ema_multi_checkpoint_synthesis(): + """Multi-checkpoint pooling reduces synthesis residual and produces distinct models. + + Tests two properties from Karras et al. 2024: + 1. More checkpoints → smaller ||w_synth - w_target|| (better approximation). + 2. Different sigma_rel targets → different synthesized model weights. + """ + import glob + + sigma_rels_train = (0.05, 0.28) + sigma_rel_target = 0.15 + num_epochs = 9 + checkpoint_every = 3 + checkpoint_epochs = list(range(checkpoint_every - 1, num_epochs, checkpoint_every)) + + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + + with tempfile.TemporaryDirectory() as tmp_dir: + train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=num_epochs, + batch_size=4, + checkpoint_every_n_epochs=checkpoint_every, + mixed_precision="no", + output_dir=tmp_dir, + force_cpu=True, + verbose=False, + ema_sigma_rels=sigma_rels_train, + ) + + assert len(glob.glob(os.path.join(tmp_dir, 'checkpoint-epoch-*'))) == 3 + + # 1. more checkpoints → smaller synthesis residual + profiles_one = compute_ema_profiles( + sigma_rels_train, checkpoint_epochs[:1], num_epochs, sigma_rel_target + ) + profiles_all = compute_ema_profiles( + sigma_rels_train, checkpoint_epochs, num_epochs, sigma_rel_target + ) + residual_one = np.sum((profiles_one['synthesized'] - profiles_one['target']) ** 2) + residual_all = np.sum((profiles_all['synthesized'] - profiles_all['target']) ** 2) + assert residual_all < residual_one, \ + f"expected residual to decrease with more checkpoints: {residual_one:.6f} → {residual_all:.6f}" + + # 2. different sigma_rel targets → different synthesized model weights + synth_lo = synthesize_ema_from_checkpoints(model, tmp_dir, sigma_rel_target=0.05) + synth_hi = synthesize_ema_from_checkpoints(model, tmp_dir, sigma_rel_target=0.28) + params_lo = torch.cat([p.flatten() for p in synth_lo.parameters()]) + params_hi = torch.cat([p.flatten() for p in synth_hi.parameters()]) + assert not torch.equal(params_lo, params_hi), \ + "sigma_rel=0.05 and sigma_rel=0.28 produced identical model weights" + + +def test_load_ema_snapshot(): + """load_ema_snapshot loads raw profile weights from a single checkpoint.""" + import glob + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + + with tempfile.TemporaryDirectory() as tmp_dir: + train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=4, batch_size=4, checkpoint_every_n_epochs=2, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + ema_sigma_rels=(0.05, 0.28), + ) + + ckpt_dirs = sorted(glob.glob(os.path.join(tmp_dir, 'checkpoint-epoch-*'))) + assert len(ckpt_dirs) == 2 + latest_ckpt = ckpt_dirs[-1] + + # different profiles should produce different weights + m0 = load_ema_snapshot(_make_unet(), latest_ckpt, profile_index=0) + m1 = load_ema_snapshot(_make_unet(), latest_ckpt, profile_index=1) + p0 = torch.cat([p.flatten() for p in m0.parameters()]) + p1 = torch.cat([p.flatten() for p in m1.parameters()]) + assert not torch.equal(p0, p1), "profile 0 and 1 snapshots are identical" + assert all(torch.isfinite(p).all() for p in m0.parameters()) + + +def test_train_ema_burn_in(): + """EMA burn-in delays update() until global_step >= ema_burn_in. + + With burn_in larger than total steps, no EMA updates fire and no .pt files + should be written. With burn_in less than total steps, snapshots exist and + their internal step counter equals total_steps - burn_in. + """ + import glob + + # case 1: burn_in exceeds total training steps → no snapshots + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + with tempfile.TemporaryDirectory() as tmp_dir: + # 2 epochs × 2 batches/epoch = 4 total steps, burn_in=10 skips them all + train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=2, batch_size=4, checkpoint_every_n_epochs=2, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + ema_sigma_rels=(0.05, 0.28), + ema_burn_in=10, + ) + pt_files = glob.glob(os.path.join(tmp_dir, 'checkpoint-epoch-*/ema/*.pt')) + assert len(pt_files) == 0, f"expected no EMA snapshots when burn_in>total, got {pt_files}" + + # case 2: burn_in skips a few steps; snapshots present and stamped with t_eff + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + with tempfile.TemporaryDirectory() as tmp_dir: + # 4 epochs × 2 batches/epoch = 8 total steps, burn_in=3 → 5 EMA updates + train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=4, batch_size=4, checkpoint_every_n_epochs=4, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + ema_sigma_rels=(0.05, 0.28), + ema_burn_in=3, + ) + pt_files = sorted(glob.glob(os.path.join(tmp_dir, 'checkpoint-epoch-*/ema/*.pt'))) + assert len(pt_files) == 2, f"expected 2 EMA snapshots, got {pt_files}" + # filename format: {profile_index}.{ema_internal_step}.pt + # ema_internal_step at the final checkpoint should be total_steps - burn_in = 8 - 3 = 5 + for f in pt_files: + stem = os.path.basename(f).replace('.pt', '') + _, step_str = stem.split('.') + assert int(step_str) == 5, f"expected internal step 5, got {step_str} from {f}" + + +def test_train_sigma_log_normal_sampling(): + """train() with sigma_log_normal samples timesteps and produces finite losses.""" + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + with tempfile.TemporaryDirectory() as tmp_dir: + result = train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10, prediction_type='v_prediction'), + num_epochs=2, batch_size=4, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + sigma_log_normal=(-1.2, 1.2), + ) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']['loss']) + + +def test_train_min_snr_weighting(): + """train() with min_snr_gamma set runs and produces finite losses for all 3 prediction types.""" + for prediction_type in ('epsilon', 'v_prediction', 'sample'): + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + with tempfile.TemporaryDirectory() as tmp_dir: + result = train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10, prediction_type=prediction_type), + num_epochs=2, batch_size=4, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + min_snr_gamma=5.0, + ) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']['loss']), \ + f"non-finite loss with min_snr_gamma + prediction_type={prediction_type}" + + +def test_train_flow_matching_basic(): + """train() with FlowMatchEulerDiscreteScheduler runs the FM noising/target + path and produces finite losses with weights that have moved.""" + model = _make_unet() + dataset = ArrayDataset(torch.randn(16, 1, 8, 8)) + + with tempfile.TemporaryDirectory() as tmp_dir: + initial_weight = model.conv_in.weight.data.clone() + result = train( + dataset, model, + noise_scheduler=FlowMatchEulerDiscreteScheduler(num_train_timesteps=10), + num_epochs=2, batch_size=4, checkpoint_every_n_epochs=2, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + ) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']['loss']) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']['epoch_loss']) + assert not torch.allclose(model.conv_in.weight.data, initial_weight) + + +def test_train_flow_matching_min_snr_is_noop(): + """min_snr_gamma is a DDPM-only weighting; setting it for FM must not error.""" + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + with tempfile.TemporaryDirectory() as tmp_dir: + result = train( + dataset, model, + noise_scheduler=FlowMatchEulerDiscreteScheduler(num_train_timesteps=10), + num_epochs=2, batch_size=4, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + min_snr_gamma=5.0, # should be silently ignored for FM + ) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']['loss']) + + +def test_train_flow_matching_sigma_log_normal(): + """sigma_log_normal is interpreted as logit-normal t parameters under FM.""" + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + with tempfile.TemporaryDirectory() as tmp_dir: + result = train( + dataset, model, + noise_scheduler=FlowMatchEulerDiscreteScheduler(num_train_timesteps=10), + num_epochs=2, batch_size=4, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + sigma_log_normal=(0.0, 1.0), + ) + assert all(torch.isfinite(torch.tensor(v)) for v in result['metrics']['loss']) + + +def test_generate_flow_matching(): + """generate() works with a FlowMatchEulerDiscreteScheduler-trained model.""" + model = _make_unet() + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=10) + images = generate(model, scheduler, batch_size=2, image_shape=(1, 8, 8), num_steps=5) + assert images.shape == (2, 1, 8, 8) + assert torch.isfinite(images).all() + + +def test_generate_s_churn_euler(): + """s_churn is forwarded to EulerDiscreteScheduler.step and produces finite output.""" + model = _make_unet() + scheduler = EulerDiscreteScheduler(num_train_timesteps=10) + images = generate( + model, scheduler, + batch_size=2, image_shape=(1, 8, 8), num_steps=5, + s_churn=10.0, s_tmin=0.0, s_tmax=float('inf'), s_noise=1.0, + ) + assert images.shape == (2, 1, 8, 8) + assert torch.isfinite(images).all() + + +def test_generate_s_churn_flow_match(): + """s_churn is also accepted by FlowMatchEulerDiscreteScheduler.""" + model = _make_unet() + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=10) + images = generate( + model, scheduler, + batch_size=2, image_shape=(1, 8, 8), num_steps=5, + s_churn=5.0, + ) + assert images.shape == (2, 1, 8, 8) + assert torch.isfinite(images).all() + + +def test_generate_s_churn_silently_dropped_for_ddpm(): + """s_churn is silently dropped for schedulers that don't accept it (no error).""" + model = _make_unet() + scheduler = DDPMScheduler(num_train_timesteps=10) + # DDPMScheduler.step doesn't accept s_churn; this should not raise. + images = generate( + model, scheduler, + batch_size=2, image_shape=(1, 8, 8), + s_churn=10.0, s_noise=2.0, + ) + assert images.shape == (2, 1, 8, 8) + assert torch.isfinite(images).all() + + +def test_train_script(): + """cosmodiff_train.py main() runs end-to-end and writes a metrics file.""" + import sys + import yaml + from scripts.cosmodiff_train import main + + minimal_config = { + "global": {"device": "cpu", "dtype": "float32"}, + "io": {"output_dir": None}, # filled in below + "data": { + "img_path": str(SIM_PATH), + "img_read_fn": "npy_read_fn", + "reshape": "2d", + "zthin": 4, + "n_samples": 8, + "keep_on_cpu": True, + "normalization": "center-max", + "norm_kwargs": {"center": None, "xmax": None, "alpha": None, "beta": None}, + }, + "augmentations": {}, + "model": { + "class": "UNet2DModel", + "kwargs": { + "sample_size": 8, + "in_channels": 1, + "out_channels": 1, + "layers_per_block": 1, + "block_out_channels": [16, 16], + "down_block_types": ["DownBlock2D", "DownBlock2D"], + "up_block_types": ["UpBlock2D", "UpBlock2D"], + "norm_num_groups": 8, + }, + }, + "noise_scheduler": { + "class": "DDPMScheduler", + "kwargs": {"num_train_timesteps": 10}, + }, + "optimizer": {"class": "AdamW", "kwargs": {"lr": 1e-4}}, + "lr_scheduler": {"class": "ConstantLR", "kwargs": {"factor": 1.0, "total_iters": 0}}, + "train": { + "num_epochs": 2, + "batch_size": 4, + "mixed_precision": "no", + "checkpoint_every_n_epochs": 2, + "force_cpu": True, + "verbose": False, + }, + } + + with tempfile.TemporaryDirectory() as tmp_dir: + minimal_config["io"]["output_dir"] = tmp_dir + config_path = os.path.join(tmp_dir, "config.yaml") + with open(config_path, "w") as f: + yaml.dump(minimal_config, f) + + orig_argv = sys.argv + try: + sys.argv = ["cosmodiff_train.py", "--config", config_path] + main() + finally: + sys.argv = orig_argv + + metrics_files = [ + f for f in os.listdir(tmp_dir) if f.startswith("metrics_epoch_") + ] + assert len(metrics_files) == 1, f"expected 1 metrics file, got {metrics_files}" + + +def test_train_script_multipath(): + """cosmodiff_train.py main() runs end-to-end on a multi-path config. + """ + import sys + import yaml + from scripts.cosmodiff_train import main + + multipath_cfg_path = DATA_PATH / 'config_multipath.yaml' + with open(multipath_cfg_path) as f: + config = yaml.safe_load(f) + + # --- override data paths and read fns for the bundled test data --- + config["data"]["img_path"] = [str(SIM_PATH), str(SIM_PATH)] + config["data"]["label_path"] = [str(PARAMS_PATH), str(PARAMS_PATH)] + config["data"]["img_read_fn"] = "npy_read_fn" + config["data"]["label_read_fn"] = "txt_read_fn" + # the shipped config still uses the legacy 'two_dim' key — replace it + config["data"].pop("two_dim", None) + config["data"]["reshape"] = "2d" + config["data"]["zthin"] = 4 + config["data"]["normalization"] = "min-max" + config["data"]["transform"] = None # log on cosmology data is fine but unrelated to this test + config["data"]["keep_on_cpu"] = True + + # --- shrink model to test scale --- + config["model"] = { + "class": "UNet2DConditionModel", + "kwargs": { + "sample_size": 64, + "in_channels": 1, + "out_channels": 1, + "layers_per_block": 1, + "block_out_channels": [16, 16], + "down_block_types": ["DownBlock2D", "DownBlock2D"], + "up_block_types": ["UpBlock2D", "UpBlock2D"], + "norm_num_groups": 8, + "cross_attention_dim": 16, + "encoder_hid_dim": N_PARAMS, + }, + } + config["noise_scheduler"] = { + "class": "DDPMScheduler", + "kwargs": {"num_train_timesteps": 10}, + } + config["augmentations"] = {} + + # --- shrink training schedule --- + config["train"].update({ + "num_epochs": 2, + "batch_size": 4, + "mixed_precision": "no", + "checkpoint_every_n_epochs": 2, + "force_cpu": True, + "verbose": False, + "conditioning": "continuous", # cosmology params are 6-D float vectors + "ema_sigma_rels": None, + "ema_burn_in": 0, + "dataloader_num_workers": 0, + }) + + with tempfile.TemporaryDirectory() as tmp_dir: + config["io"]["output_dir"] = tmp_dir + config_path = os.path.join(tmp_dir, "config_multipath.yaml") + with open(config_path, "w") as f: + yaml.dump(config, f) + + orig_argv = sys.argv + try: + sys.argv = ["cosmodiff_train.py", "--config", config_path] + main() + finally: + sys.argv = orig_argv + + metrics_files = [ + f for f in os.listdir(tmp_dir) if f.startswith("metrics_epoch_") + ] + assert len(metrics_files) == 1, f"expected 1 metrics file, got {metrics_files}" + + +def test_sample_script(): + """cosmodiff_sample.py main() runs end-to-end and writes an .npy file.""" + import sys + from scripts.cosmodiff_sample import main + + # First train a tiny model so we have a real checkpoint on disk. + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + + with tempfile.TemporaryDirectory() as tmp_dir: + train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=2, batch_size=4, checkpoint_every_n_epochs=2, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + ) + + out_npy = os.path.join(tmp_dir, "samples.npy") + orig_argv = sys.argv + try: + sys.argv = [ + "cosmodiff_sample.py", + "--output_dir", tmp_dir, + "--n_samples", "4", + "--batch_size", "2", + "--image_shape", "1", "8", "8", + "--scheduler", "DDIMScheduler", + "--num_steps", "5", + "--device", "cpu", + "--output", out_npy, + "--seed", "0", + ] + main() + finally: + sys.argv = orig_argv + + assert os.path.exists(out_npy), f"expected samples at {out_npy}" + arr = np.load(out_npy) + assert arr.shape == (4, 1, 8, 8), f"unexpected shape {arr.shape}" + assert np.isfinite(arr).all(), "non-finite values in generated samples" + + +def test_sample_script_with_ema(): + """cosmodiff_sample.py with --ema_sigma_rel synthesizes EMA before sampling.""" + import sys + from scripts.cosmodiff_sample import main + + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + + with tempfile.TemporaryDirectory() as tmp_dir: + train( + dataset, model, + noise_scheduler=DDPMScheduler(num_train_timesteps=10), + num_epochs=4, batch_size=4, checkpoint_every_n_epochs=2, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + ema_sigma_rels=(0.05, 0.28), + ) + + out_npy = os.path.join(tmp_dir, "samples_ema.npy") + orig_argv = sys.argv + try: + sys.argv = [ + "cosmodiff_sample.py", + "--output_dir", tmp_dir, + "--n_samples", "2", + "--image_shape", "1", "8", "8", + "--scheduler", "DDIMScheduler", + "--num_steps", "2", + "--device", "cpu", + "--output", out_npy, + "--ema_sigma_rel", "0.05", + "--seed", "0", + ] + main() + finally: + sys.argv = orig_argv + + assert os.path.exists(out_npy) + arr = np.load(out_npy) + assert arr.shape == (2, 1, 8, 8) + assert np.isfinite(arr).all() + + +def test_sample_script_flow_matching(): + """cosmodiff_sample.py runs end-to-end against a FM-trained checkpoint.""" + import sys + from scripts.cosmodiff_sample import main + + model = _make_unet() + dataset = ArrayDataset(torch.randn(8, 1, 8, 8)) + + with tempfile.TemporaryDirectory() as tmp_dir: + train( + dataset, model, + noise_scheduler=FlowMatchEulerDiscreteScheduler(num_train_timesteps=10), + num_epochs=2, batch_size=4, checkpoint_every_n_epochs=2, + mixed_precision="no", output_dir=tmp_dir, force_cpu=True, verbose=False, + ) + + out_npy = os.path.join(tmp_dir, "samples_fm.npy") + orig_argv = sys.argv + try: + sys.argv = [ + "cosmodiff_sample.py", + "--output_dir", tmp_dir, + "--n_samples", "2", + "--image_shape", "1", "8", "8", + "--num_steps", "5", + "--device", "cpu", + "--output", out_npy, + "--seed", "0", + ] + main() + finally: + sys.argv = orig_argv + assert os.path.exists(out_npy) + arr = np.load(out_npy) + assert arr.shape == (2, 1, 8, 8) + assert np.isfinite(arr).all() diff --git a/cosmodiff/tests/test_transform.py b/cosmodiff/tests/test_transform.py new file mode 100644 index 0000000..003075b --- /dev/null +++ b/cosmodiff/tests/test_transform.py @@ -0,0 +1,225 @@ +import copy +import numpy as np +import torch +from cosmodiff.transform import ( + minmax_norm, + center_max_norm, + tanh_norm, + Normalization, + Transform, +) + + +# --------------------------------------------------------------------------- +# minmax_norm / center_max_norm / tanh_norm +# --------------------------------------------------------------------------- + +def test_minmax_norm(): + # use non-negative data so the formula x *= 2/xmax is well-defined + x = torch.rand(4, 8, 8) + out, params = minmax_norm(x) + assert out.min().item() >= -1.0 - 1e-6 + assert out.max().item() <= 1.0 + 1e-6 + assert 'xmin' in params and 'xmax' in params + + +def test_minmax_norm_inverse(): + x = torch.rand(4, 8, 8) + out, params = minmax_norm(x) + recovered, _ = minmax_norm(out, inverse=True, **params) + assert torch.allclose(recovered, x, atol=1e-5) + + +def test_minmax_norm_inplace(): + x = torch.rand(4, 8, 8) + x_clone = x.clone() + out, _ = minmax_norm(x, inplace=True) + assert out.data_ptr() == x.data_ptr() + _, params2 = minmax_norm(x_clone) + assert torch.allclose(out, x_clone.sub(params2['xmin']).mul(2 / params2['xmax']).sub(1), atol=1e-5) + + +def test_center_max_norm(): + x = torch.randn(100) + out, params = center_max_norm(x.clone()) + assert abs(out.mean().item()) < 0.1 + assert out.abs().max().item() <= 1.0 + 1e-6 + assert 'center' in params and 'xmax' in params + + +def test_center_max_norm_inverse(): + x = torch.randn(100) + out, params = center_max_norm(x.clone()) + recovered, _ = center_max_norm(out, inverse=True, **params) + assert torch.allclose(recovered, x, atol=1e-5) + + +def test_tanh_norm(): + x = torch.randn(100) + out, params = tanh_norm(x) + assert out.shape == x.shape + assert set(params.keys()) == {'mu', 'alpha', 'beta', 'delta', 'gamma', 'sigma'} + # tanh output is strictly bounded by sigma * (-beta, alpha) + assert out.min().item() > -params['beta'] * params['sigma'] + assert out.max().item() < params['alpha'] * params['sigma'] + + +def test_tanh_norm_sigma(): + x = torch.randn(100) + out_default, _ = tanh_norm(x, sigma=1.0) + out_scaled, params = tanh_norm(x, sigma=2.0) + assert torch.allclose(out_scaled, out_default * 2.0, atol=1e-6) + assert params['sigma'] == 2.0 + + +def test_tanh_norm_inverse(): + x = torch.randn(100) * 0.5 # keep within tanh saturation limits + out, params = tanh_norm(x) + recovered, _ = tanh_norm(out, inverse=True, **params) + assert torch.allclose(recovered, x, atol=1e-5) + + +# --------------------------------------------------------------------------- +# Normalization class +# --------------------------------------------------------------------------- + +def test_normalization_minmax(): + x = torch.rand(10, 1, 8, 8) + norm = Normalization('min-max', inplace=False) + out = norm(x) + assert out is not None + assert out.shape == x.shape + assert out.min().item() >= -1.0 - 1e-6 + assert out.max().item() <= 1.0 + 1e-6 + + +def test_normalization_centermax(): + x = torch.randn(10, 1, 8, 8) + norm = Normalization('center-max', inplace=False) + out = norm(x) + assert out is not None + assert out.shape == x.shape + assert out.abs().max().item() <= 1.0 + 1e-6 + + +def test_normalization_inverse(): + x = torch.rand(10, 1, 8, 8) + norm = Normalization('min-max', inplace=False) + out = norm(x) + recovered = norm.inverse(out) + assert torch.allclose(recovered, x, atol=1e-5) + + +def test_normalization_tanh(): + # tanh branch: minmax first, then tanh — output is strictly within sigma*(-beta, alpha) + x = torch.rand(10, 1, 8, 8) + norm = Normalization('tanh', inplace=False) + out = norm(x) + assert out.shape == x.shape + assert out.min().item() > -norm.kwargs['beta'] * norm.kwargs['sigma'] + assert out.max().item() < norm.kwargs['alpha'] * norm.kwargs['sigma'] + # params from both stages must be stored after forward + assert all(k in norm.kwargs for k in ('center', 'xmax', 'mu', 'alpha', 'beta', 'gamma', 'delta', 'sigma')) + + +def test_normalization_tanh_inverse(): + x = torch.rand(10, 1, 8, 8) + norm = Normalization('tanh', inplace=False) + out = norm(x) + recovered = norm.inverse(out) + assert torch.allclose(recovered, x, atol=1e-5) + + +def test_normalization_params_fixed_when_given(): + """Params supplied at init must not be overwritten by the forward pass.""" + fixed_xmin = torch.tensor(0.0) + fixed_xmax = torch.tensor(2.0) + norm = Normalization('min-max', inplace=False, xmin=fixed_xmin, xmax=fixed_xmax) + + # forward on data whose natural min/max differ from the fixed values + x = torch.rand(10, 1, 8, 8) * 10 + 5 + norm(x) + + assert torch.equal(norm.kwargs['xmin'], fixed_xmin), "xmin was overwritten" + assert torch.equal(norm.kwargs['xmax'], fixed_xmax), "xmax was overwritten" + + +def test_normalization_params_inferred_on_first_forward(): + """Params not supplied at init must be populated after the first forward pass.""" + norm = Normalization('min-max', inplace=False) + assert 'xmin' not in norm.kwargs + assert 'xmax' not in norm.kwargs + + x = torch.rand(10, 1, 8, 8) + norm(x) + + assert 'xmin' in norm.kwargs, "xmin not set after first forward" + assert 'xmax' in norm.kwargs, "xmax not set after first forward" + assert np.allclose(norm.kwargs['xmin'], x.min().item()) + assert np.allclose(norm.kwargs['xmax'], x.max().item()) + + # second forward on different data must reuse the inferred params, not recompute + y = torch.rand(10, 1, 8, 8) * 10 + 5 + xmin_after_first = copy.copy(norm.kwargs['xmin']) + xmax_after_first = copy.copy(norm.kwargs['xmax']) + norm(y) + + assert np.allclose(norm.kwargs['xmin'], xmin_after_first), "xmin changed on second forward" + assert np.allclose(norm.kwargs['xmax'], xmax_after_first), "xmax changed on second forward" + + +# --------------------------------------------------------------------------- +# Transform class +# --------------------------------------------------------------------------- + +def test_transform_log_roundtrip(): + """log forward + inverse recovers the original tensor.""" + t = Transform(ops=[], log=True) + x = torch.rand(2, 1, 8, 8) + 0.5 # strictly positive + y = t(x) + assert torch.allclose(y, x.log(), atol=1e-6) + x_back = t.inverse(y) + assert torch.allclose(x_back, x, atol=1e-5) + + +def test_transform_fft2_shape_and_roundtrip(): + """fft2 op doubles the channel dim (real|imag) and round-trips back.""" + t = Transform(ops=['fft2'], log=False) + x = torch.randn(3, 2, 8, 8) + y = t(x) + # real and imaginary parts concatenated along channel dim + assert y.shape == (3, 4, 8, 8), f"expected channel doubling, got {y.shape}" + x_back = t.inverse(y) + assert x_back.shape == x.shape + assert torch.allclose(x_back, x, atol=1e-5) + + +def test_transform_log_then_fft2_roundtrip(): + """log → fft2 forward; ifft2 → exp inverse round-trips.""" + t = Transform(ops=['fft2'], log=True) + x = torch.rand(2, 1, 8, 8) + 0.5 # strictly positive so log is well-defined + y = t(x) + assert y.shape == (2, 2, 8, 8) + x_back = t.inverse(y) + assert torch.allclose(x_back, x, atol=1e-5) + + +def test_transform_inplace_false_does_not_mutate_input(): + """Default inplace=False leaves the input tensor untouched.""" + t = Transform(ops=[], log=True) + x = torch.rand(2, 1, 8, 8) + 0.5 + x_orig = x.clone() + _ = t(x) + assert torch.equal(x, x_orig), "input was mutated despite inplace=False" + + +def test_transform_unknown_op_raises(): + """Unknown op names raise a clear ValueError.""" + t = Transform(ops=['not-a-real-op']) + x = torch.randn(2, 1, 8, 8) + try: + t(x) + except ValueError as e: + assert "not-a-real-op" in str(e) + else: + raise AssertionError("expected ValueError for unknown op") diff --git a/cosmodiff/tests/test_utils.py b/cosmodiff/tests/test_utils.py index 11af6f5..4897798 100644 --- a/cosmodiff/tests/test_utils.py +++ b/cosmodiff/tests/test_utils.py @@ -1,3 +1,4 @@ +import copy import json import tempfile import numpy as np @@ -5,14 +6,22 @@ from cosmodiff.utils import ( ArrayDataset, load_data, - minmax_norm, - center_scale_norm, npy_read_fn, + txt_read_fn, parse_config_model, parse_config_data, + read_config, write_metrics, read_metrics, ) +from cosmodiff.transform import Normalization, MultiNormalization, Transform, MultiTransform +from cosmodiff.data import DATA_PATH + +CONFIG_PATH = DATA_PATH / 'config.yaml' +SIM_PATH = DATA_PATH / 'IllustrisTNG_Mcdm.npy' +PARAMS_PATH = DATA_PATH / 'params_Illustris.txt' +# shape: (34, 32, 32, 32), dtype: float32 +DATA_N, DATA_NZ, DATA_NX, DATA_NY = 34, 32, 32, 32 def _make_array(n=20, nz=4, nx=8, ny=8): @@ -24,16 +33,14 @@ def _make_array(n=20, nz=4, nx=8, ny=8): # --------------------------------------------------------------------------- def test_n_samples(): - arr = _make_array(n=20) - images, _ = load_data(arr, img_read_fn=None, norm=None, two_dim=False) - assert images.shape[0] == 20 + images_all = load_data(SIM_PATH, img_read_fn=npy_read_fn, normalization=None, reshape="3d")['images'] + assert images_all.shape[0] == DATA_N - images, _ = load_data(arr, img_read_fn=None, n_samples=7, norm=None, two_dim=False) - assert images.shape[0] == 7 + images_sub = load_data(SIM_PATH, img_read_fn=npy_read_fn, n_samples=3, normalization=None, reshape="3d")['images'] + assert images_sub.shape[0] == 3 - arr_t = torch.as_tensor(arr) - for i in range(len(images)): - assert any(torch.allclose(images[i], arr_t[j]) for j in range(len(arr_t))) + for i in range(len(images_sub)): + assert any(torch.allclose(images_sub[i], images_all[j]) for j in range(len(images_all))) def test_n_samples_labels_in_sync(): @@ -43,12 +50,13 @@ def test_n_samples_labels_in_sync(): # labels encode each sample's original row index so we can verify alignment labels = np.arange(20) - images, out_labels = load_data( + out = load_data( arr, img_read_fn=None, label_path=labels, label_read_fn=None, n_samples=7, seed=0, - norm=None, two_dim=False, + normalization=None, reshape="3d", ) + images, out_labels = out['images'], out['labels'] assert images.shape[0] == 7 assert out_labels.shape[0] == 7 @@ -59,23 +67,110 @@ def test_n_samples_labels_in_sync(): def test_seed(): - arr = _make_array(n=50) - imgs1, _ = load_data(arr, img_read_fn=None, n_samples=10, seed=42, norm=None, two_dim=False) - imgs2, _ = load_data(arr, img_read_fn=None, n_samples=10, seed=42, norm=None, two_dim=False) - imgs3, _ = load_data(arr, img_read_fn=None, n_samples=10, seed=99, norm=None, two_dim=False) + imgs1 = load_data(SIM_PATH, img_read_fn=npy_read_fn, n_samples=3, seed=42, normalization=None, reshape="3d")['images'] + imgs2 = load_data(SIM_PATH, img_read_fn=npy_read_fn, n_samples=3, seed=42, normalization=None, reshape="3d")['images'] + imgs3 = load_data(SIM_PATH, img_read_fn=npy_read_fn, n_samples=3, seed=99, normalization=None, reshape="3d")['images'] assert torch.allclose(imgs1, imgs2) assert not torch.allclose(imgs1, imgs3) def test_memmap(): - arr = _make_array(n=10) - with tempfile.NamedTemporaryFile(suffix=".npy") as f: - np.save(f.name, arr) - mmap = np.load(f.name, mmap_mode="r") - images_all, _ = load_data(mmap, img_read_fn=None, norm=None, two_dim=False) - images_sub, _ = load_data(mmap, img_read_fn=None, n_samples=5, norm=None, two_dim=False) - assert images_all.shape[0] == 10 - assert images_sub.shape[0] == 5 + mmap = np.load(SIM_PATH, mmap_mode="r") + images_all = load_data(mmap, img_read_fn=None, normalization=None, reshape="3d")['images'] + images_sub = load_data(mmap, img_read_fn=None, n_samples=3, normalization=None, reshape="3d")['images'] + assert images_all.shape[0] == DATA_N + assert images_sub.shape[0] == 3 + + +# --------------------------------------------------------------------------- +# load_data: multi-path returns the right (Multi)Normalization / (Multi)Transform +# --------------------------------------------------------------------------- + +def test_load_data_multipath_norm_transform_types(): + """Two img_paths + two label_paths exercising the two normalization modes. + + Mode A — single normalization / transform shared across all paths + → ``out['norm']`` is :class:`Normalization` + → ``out['tform']`` is :class:`Transform` + + Mode B — list-shaped normalization / transform (one per path) + → ``out['norm']`` is :class:`MultiNormalization` + → ``out['tform']`` is :class:`MultiTransform` + """ + img_paths = [str(SIM_PATH), str(SIM_PATH)] + label_paths = [str(PARAMS_PATH), str(PARAMS_PATH)] + + # --- Case (i): scalar normalization / transform --- + out_a = load_data( + img_path=img_paths, + img_read_fn=npy_read_fn, + label_path=label_paths, + label_read_fn=txt_read_fn, + reshape='2d', + zthin=4, + normalization='min-max', + transform=['log'], + ) + assert isinstance(out_a['norm'], Normalization) + assert not isinstance(out_a['norm'], MultiNormalization) + assert isinstance(out_a['tform'], Transform) + assert not isinstance(out_a['tform'], MultiTransform) + + # --- Case (ii): per-path normalization / transform --- + out_b = load_data( + img_path=img_paths, + img_read_fn=npy_read_fn, + label_path=label_paths, + label_read_fn=txt_read_fn, + reshape='2d', + zthin=4, + normalization=['min-max', 'min-max'], + norm_kwargs=[{}, {}], + transform=[['log'], ['log']], + ) + assert isinstance(out_b['norm'], MultiNormalization) + assert isinstance(out_b['tform'], MultiTransform) + + +def test_load_data_log_fft2_transform_roundtrip(): + """``transform=['log', 'fft2']`` via load_data round-trips back to the + raw input under tform.inverse. + + Exercises the full config-style path: ``load_data`` parses the list, + extracts ``log`` as a flag, constructs the :class:`Transform`, applies + it to the loaded images, and stores the fitted object in ``out['tform']``. + Calling ``.inverse()`` on the loaded images recovers the post-reshape + raw images within floating-point tolerance. + """ + # synthetic O(1)-magnitude positive data so log() and the FFT round-trip + # don't accumulate magnitude-related FP error (cosmology data sits near + # 1e10 which exceeds float32 round-trip precision) + rng = np.random.default_rng(0) + arr = rng.random((4, 8, 16, 16)).astype(np.float32) + 0.1 + + out_raw = load_data( + img_path=arr, + img_read_fn=None, + reshape='2d', + zthin=2, + normalization=None, + transform=None, + ) + out_tf = load_data( + img_path=arr, + img_read_fn=None, + reshape='2d', + zthin=2, + normalization=None, + transform=['log', 'fft2'], + ) + assert isinstance(out_tf['tform'], Transform) + # forward doubled the channel dim via fft2 real|imag concat + assert out_tf['images'].shape[1] == out_raw['images'].shape[1] * 2 + # inverse recovers the raw input + recovered = out_tf['tform'].inverse(out_tf['images']) + assert torch.allclose(recovered, out_raw['images'], atol=1e-5), \ + f"max diff after log→fft2 roundtrip: {(recovered - out_raw['images']).abs().max()}" # --------------------------------------------------------------------------- @@ -96,35 +191,14 @@ def test_array_dataset(): assert torch.allclose(item["images"], images[3]) -# --------------------------------------------------------------------------- -# minmax_norm / center_scale_norm -# --------------------------------------------------------------------------- - -def test_minmax_norm(): - x = torch.randn(4, 8, 8) - out = minmax_norm(x) - assert out.min().item() >= -1.0 - 1e-6 - assert out.max().item() <= 1.0 + 1e-6 - - -def test_center_scale_norm(): - x = torch.randn(100) - out = center_scale_norm(x.clone()) - assert abs(out.mean().item()) < 0.1 - assert out.abs().max().item() <= 1.0 + 1e-6 - - # --------------------------------------------------------------------------- # npy_read_fn # --------------------------------------------------------------------------- def test_npy_read_fn(): - arr = np.random.rand(5, 4, 8, 8).astype(np.float32) - with tempfile.NamedTemporaryFile(suffix=".npy") as f: - np.save(f.name, arr) - result = npy_read_fn(f.name) - assert result.shape == arr.shape - assert np.allclose(result, arr) + result = npy_read_fn(SIM_PATH) + assert result.shape == (DATA_N, DATA_NZ, DATA_NX, DATA_NY) + assert result.dtype == np.float32 # --------------------------------------------------------------------------- @@ -163,11 +237,11 @@ def test_parse_config_model(): "noise_scheduler": {"class": "DDPMScheduler", "kwargs": {"num_train_timesteps": 10}}, "lr_scheduler": {"class": "ConstantLR", "kwargs": {"factor": 1.0, "total_iters": 0}}, } - model, optimizer, noise_scheduler, lr_scheduler = parse_config_model(config) - assert model is not None - assert optimizer is not None - assert noise_scheduler is not None - assert lr_scheduler is not None + out = parse_config_model(config) + assert out['model'] is not None + assert out['optimizer'] is not None + assert out['noise_scheduler'] is not None + assert out['lr_scheduler'] is not None # --------------------------------------------------------------------------- @@ -185,15 +259,79 @@ def test_parse_config_data(): "data": { "img_path": img_path, "img_read_fn": "npy_read_fn", - "log": False, - "norm": "min-max", - "two_dim": True, + "normalization": "min-max", + "reshape": "2d", "zthin": 1, "keep_on_cpu": True, }, } - dataset = parse_config_data(config) + out = parse_config_data(config) + dataset = out['data'] assert len(dataset) == 6 * 4 # Nbatch * Nz slices item = dataset[0] assert "images" in item assert item["images"].shape == torch.Size([1, 8, 8]) + + +# --------------------------------------------------------------------------- +# configs/config.yaml round-trip +# --------------------------------------------------------------------------- + +def test_config_yaml_parse_model(): + """parse_config_model must instantiate all four components from config.yaml.""" + config = read_config(CONFIG_PATH) + config['global']['device'] = 'cpu' # don't require GPU in CI + + out = parse_config_model(config) + model = out['model'] + optimizer = out['optimizer'] + noise_scheduler = out['noise_scheduler'] + lr_scheduler = out['lr_scheduler'] + + assert type(model).__name__ == 'UNet2DModel' + assert model.config.sample_size == 64 + assert model.config.in_channels == 1 + assert model.config.out_channels == 1 + + assert type(optimizer).__name__ == 'AdamW' + assert abs(optimizer.param_groups[0]['lr'] - 1e-4) < 1e-10 + assert abs(optimizer.param_groups[0]['weight_decay'] - 1e-2) < 1e-10 + + assert type(noise_scheduler).__name__ == 'DDPMScheduler' + assert noise_scheduler.config.num_train_timesteps == 1000 + + assert type(lr_scheduler).__name__ == 'ConstantLR' + + +def test_config_yaml_parse_data(): + """parse_config_data must build an ArrayDataset with the right shape, + normalization, and augmentation pipeline from config.yaml.""" + from cosmodiff.augment import RandomRoll, RandomFlip + + config = read_config(CONFIG_PATH) + config['global']['device'] = 'cpu' + + cfg = copy.deepcopy(config) + cfg['data']['img_path'] = str(SIM_PATH) + cfg['data']['label_path'] = None + cfg['data']['keep_on_cpu'] = True + cfg['data']['transform'] = None # this test checks raw shape, not transform behavior + + out = parse_config_data(cfg) + dataset = out['data'] + norm = out['norm'] + + zthin = config['data']['zthin'] + assert isinstance(dataset, ArrayDataset) + assert len(dataset) == DATA_N * (DATA_NZ // zthin) + item = dataset[0] + assert 'images' in item + assert item['images'].shape == torch.Size([1, DATA_NX, DATA_NY]) + + assert isinstance(norm, Normalization) + assert norm.method == config['data']['normalization'] + + assert dataset.augmentations is not None + aug_types = [type(t) for t in dataset.augmentations] + assert RandomRoll in aug_types + assert RandomFlip in aug_types diff --git a/cosmodiff/transform.py b/cosmodiff/transform.py new file mode 100644 index 0000000..03fa458 --- /dev/null +++ b/cosmodiff/transform.py @@ -0,0 +1,333 @@ +"""Data normalization transforms for cosmodiff. + +Provides three pointwise normalization functions and a :class:`Normalization` +``nn.Module`` that wraps them with a uniform forward / inverse interface. +""" + +import torch + + +def minmax_norm( + x: torch.Tensor, + xmin: float | None = None, + xmax: float | None = None, + inverse: bool = False, + inplace: bool = False, + **kwargs +) -> torch.Tensor: + """Normalize a tensor to ``[-1, 1]`` via min-max scaling. + + Args: + x (torch.Tensor): Input tensor of any shape. + xmin (float): normalize by min, default is x.min() + xmax (float): normalize by max, default is x.max() + inverse (bool): if True, apply the inverse mapping (requires xmin, xmax) + inplace (bool): edit tensor inplace, default is False + + Returns: + torch.Tensor: Normalized tensor with values in ``[-1, 1]``. + dict: normalization parameters + """ + if not inplace: + x = x.clone() + if not inverse: + if xmin is None: + xmin = x.min().item() + if xmax is None: + xmax = x.max().item() + x -= xmin + x *= 2 / xmax + x -= 1 + + else: + assert xmin is not None + assert xmax is not None + x = (x + 1) / 2 * xmax + xmin + + return x, {'xmin': xmin, 'xmax': xmax} + + +def center_max_norm( + x: torch.Tensor, + center: float | None = None, + xmax: float | None = None, + inverse: bool | None = False, + inplace: bool = False, + **kwargs +): + """Center a tensor based on its average, and normalize by its absolute deviation. + + Args: + x (torch.Tensor): Input tensor of any shape. + center (float): average centering to subtract off + xmax (float): abs-max normalization to divide by + inverse (bool): apply inverse operation, requires center and xmax + inplace (bool): If True, edit inplace. + + Returns: + torch.Tensor: scaled tensor + dict: normalization parameters + """ + if not inplace: + x = x.clone() + + if not inverse: + # center + if center is None: + center = x.mean().item() + x -= center + + # scale by max-abs + if xmax is None: + xmax = x.abs().max().item() + x /= xmax + + else: + assert xmax is not None + assert center is not None + x *= xmax + x += center + + return x, {'center': center, 'xmax': xmax} + + +def tanh_norm(x, alpha=1.0, beta=1.0, gamma=1.0, delta=1.0, sigma=1.0, mu=0.0, inverse=False, **kwargs): + """ + tanh normalization + + f(x) = sigma * alpha * tanh((gamma * (x-mu)) / alpha) if (x-mu) >= 0 + f(x) = sigma * beta * tanh((delta * (x-mu)) / beta) if (x-mu) < 0 + + Args: + x (tensor): Input tensor + alpha (float): Positive saturation limit (upper bound). + beta (float): Negative saturation limit (lower bound). + gamma (float): Multiplicative gain for the positive side. + delta (float): Multiplicative gain for the negative side. + sigma (float): Final scaling + mu (float): Mean shift / center point. + inverse (bool): Toggle between forward and inverse operations. + + Returns: + torch.Tensor: normalized data + dict: normalization parameters + """ + if not inverse: + # Forward: x -> y + x_shifted = x - mu + pos = alpha * torch.tanh((gamma * x_shifted) / alpha) + neg = beta * torch.tanh((delta * x_shifted) / beta) + y = torch.where(x_shifted >= 0, pos, neg) * sigma + else: + # Inverse: y -> x + # Note: data (y) should be clamped within (-beta, alpha) for stability + y = torch.clamp(x / sigma, -beta + 1e-9, alpha - 1e-9) + pos_inv = (alpha * torch.atanh(y / alpha)) / gamma + neg_inv = (beta * torch.atanh(y / beta)) / delta + y = torch.where(y >= 0, pos_inv, neg_inv) + mu + + params = {'mu': mu, 'alpha': alpha, 'beta': beta, 'delta': delta, 'gamma': gamma, 'sigma': sigma} + return y, params + + +class Normalization(torch.nn.Module): + """Invertible pointwise data normalization. + + Wraps :func:`minmax_norm`, :func:`center_max_norm`, and :func:`tanh_norm` + behind a single object with consistent ``forward`` / ``inverse`` calls. + On the first forward pass, any normalization parameters not supplied at + construction (e.g. ``xmin``, ``xmax``, ``center``) are estimated from the + input tensor and stored on ``self.kwargs``; subsequent forward and inverse + calls reuse those fitted parameters so the mapping stays consistent + across batches. + + Args: + method (str): which normalization to apply. One of: + + * ``'minmax'`` / ``'min-max'``: linear rescale to ``[-1, 1]``. + * ``'centermax'`` / ``'center-max'``: subtract mean, divide by + max-abs. + * ``'tanh'``: ``center-max`` followed by :func:`tanh_norm` for + soft saturation of heavy tails. + + inplace (bool): edit the tensor in place when supported. Defaults to + ``False`` (returns a new tensor). + **kwargs: parameters forwarded to the underlying normalization + function (e.g. ``xmin``, ``xmax``, ``center``, ``alpha``, + ``beta``, ``gamma``, ``delta``, ``sigma``, ``mu``). Any value + left as ``None`` is estimated from data on the first forward pass + and persisted on ``self.kwargs``. + + Example:: + + norm = Normalization('center-max') + y = norm(x) # fit on x, return normalized + x_back = norm.inverse(y) # exactly recovers x + """ + + def __init__(self, method: str, inplace: bool = False, **kwargs): + super().__init__() + self.method = method + self.inplace = inplace + self.kwargs = kwargs + + def forward(self, x): + if self.method in ['minmax', 'min-max']: + x, kw = minmax_norm(x, inplace=self.inplace, **self.kwargs) + self.kwargs.update(kw) + + elif self.method in ['centermax', 'center-max']: + x, kw = center_max_norm(x, inplace=self.inplace, **self.kwargs) + self.kwargs.update(kw) + + elif self.method in ['tanh']: + x, kw = center_max_norm(x, inplace=self.inplace, **self.kwargs) + self.kwargs.update(kw) + x, kw = tanh_norm(x, inplace=self.inplace, **self.kwargs) + self.kwargs.update(kw) + + return x + + def inverse(self, x): + if self.method in ['minmax', 'min-max']: + x, _ = minmax_norm(x, inplace=self.inplace, inverse=True, **self.kwargs) + elif self.method in ['centermax', 'center-max']: + x, _ = center_max_norm(x, inplace=self.inplace, inverse=True, **self.kwargs) + elif self.method in ['tanh']: + x, _ = tanh_norm(x, inplace=self.inplace, inverse=True, **self.kwargs) + x, _ = center_max_norm(x, inplace=self.inplace, inverse=True, **self.kwargs) + + return x + + def __repr__(self): + return f"{self.__class__.__name__}(method={self.method})" + + +class MultiNormalization(torch.nn.Module): + """ + Multiple normalizations applied to batch dimension + of data for a set of discrete class labels. + WIP + """ + + def __init__(self, classes, norms): + self.classes = classes + self.norms = norms + + def forward(self, x, labels): + for lbl, norm in zip(self.classes, self.norms): + idx = labels == lbl + x[idx] = norm(x[idx]) + + return x + + def inverse(self, x, labels): + for lbl, norm in zip(self.classes, self.norms): + idx = labels == lbl + x[idx] = norm.inverse(x[idx]) + + return x + + +class Transform(torch.nn.Module): + """Composable, invertible data transformation pipeline. + + Applies an optional ``log`` transform first, then iterates through a list + of named operations in ``ops``. Each operation must have a defined + inverse; calling :meth:`inverse` runs the operations in reverse order. + + Args: + ops (list of str, optional): ordered list of operation names to apply + after the log step. Currently supported: + + * ``'fft2'``: 2D FFT; real and imaginary parts are concatenated + along the channel dimension to keep the tensor real-valued. + Actually uses fft2 so that we don't change size of image. + But negative frequencies are discarded upon ifft2. + + Defaults to an empty pipeline if ``None``. + log (bool): if ``True``, take ``log()`` before any ``ops`` (and + ``exp()`` last on inverse). Useful for log-normal data such as + cosmological matter density fields. Defaults to ``False``. + inplace (bool): edit the tensor in place when supported. Defaults + to ``False``. + """ + def __init__( + self, + ops: list[str] | None = None, + log: bool = False, + inplace: bool = False, + **kwargs, + ): + super().__init__() + self.ops = ops + self.log = log + self.inplace = inplace + + def forward(self, x): + if not self.inplace: + x = x.clone() + + # first operation is log + if self.log: + x.log_() + + # now perform other operations + for op in self.ops: + if op == 'fft2': + x = torch.fft.fft2(x) + x = torch.cat([x.real, x.imag], dim=1) + + else: + raise ValueError(f"didn't recognize '{op}'") + + return x + + def inverse(self, x): + if not self.inplace: + x = x.clone() + + # go through operations in reverse order + for op in self.ops[::-1]: + if op == 'fft2': + N = x.shape[1]//2 + x = torch.complex(x[:, :N], x[:, N:]) + x = torch.fft.ifft2(x).real + + else: + raise ValueError(f"didn't recognize '{op}'") + + # last operation is unlog + if self.log: + x.exp_() + + return x + + def __repr__(self): + return f"{self.__class__.__name__}(log={self.log}, ops={self.ops})" + + +class MultiTransform(torch.nn.Module): + """ + Multiple transformations applied to batch dimension + of data for a set of discrete class labels. + WIP + """ + def __init__(self, classes, tforms): + self.classes = classes + self.tforms = tforms + + def forward(self, x, labels): + for lbl, tform in zip(self.classes, self.tforms): + idx = labels == lbl + x[idx] = tform(x[idx]) + + return x + + def inverse(self, x): + for lbl, tform in zip(self.classes, self.tforms): + idx = labels == lbl + x[idx] = tform.inverse(x[idx]) + + return x + diff --git a/cosmodiff/utils.py b/cosmodiff/utils.py index 050fb88..cb79b7b 100644 --- a/cosmodiff/utils.py +++ b/cosmodiff/utils.py @@ -14,7 +14,9 @@ from tqdm.auto import tqdm import time import yaml -from typing import Optional +from typing import Optional, Callable + +from .transform import Normalization, MultiNormalization, Transform, MultiTransform class ArrayDataset(Dataset): @@ -41,11 +43,12 @@ class ArrayDataset(Dataset): dataset = ArrayDataset(images, labels=labels) # {"images": tensor, "labels": tensor} """ - def __init__(self, arrays: torch.Tensor, labels: torch.Tensor = None, augmentations: callable = None): - # CUDA tensors cannot be shared across DataLoader worker processes (fork - # cannot inherit the CUDA context). Move to CPU; the training loop - # (via accelerate) handles GPU placement per-batch. - self.arrays = arrays.cpu() + def __init__(self, + arrays: torch.Tensor, + labels: torch.Tensor = None, + augmentations: callable = None, + ): + self.arrays = arrays self.labels = labels self.augmentations = augmentations @@ -62,18 +65,20 @@ def __getitem__(self, idx): def load_data( - img_path: str | np.ndarray, - img_read_fn: callable, + img_path: str | np.ndarray | list[str], + img_read_fn: Callable | list[Callable], device: str | None = None, dtype: torch.dtype | None = None, - label_path: str | np.ndarray | None = None, - label_read_fn: Optional[callable] = None, - log: bool = False, - norm: str | None = None, - two_dim: bool = True, - zthin: int = 1, - n_samples: int | None = None, + label_path: str | np.ndarray | list[str] | None = None, + label_read_fn: Optional[Callable | list[Callable]] = None, + reshape: str | None = '2d', + zthin: int | list[int] = 1, + n_samples: int | list[int] | None = None, seed: np.random.Generator | None = None, + normalization: str | list[str] | None = None, + norm_kwargs: dict | list[dict] | None = None, + transform: list[str] | list[list[str]] | None = None, + read_only: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Load images and optionally labels into tensors ready for ``ArrayDataset``. @@ -83,14 +88,6 @@ def load_data( Input arrays are assumed to be of shape ``(Nbatch, Nz, Nx, Ny)``. - If ``two_dim=True``, the z axis is optionally thinned by ``zthin`` and - the array is reshaped to ``(Nbatch * Nz, 1, Nx, Ny)``, treating each - z-slice as an independent 2D image. - - If ``two_dim=False``, ``zthin`` is ignored and the array is unsqueezed - to ``(Nbatch, 1, Nz, Nx, Ny)`` treating each sample as a 3D volume - with a single channel. - Args: img_path (str or np.ndarray): Path to the image data file on disk, or a pre-loaded numpy array. @@ -107,21 +104,30 @@ def load_data( label_read_fn (callable, optional): User-provided function that accepts ``label_path`` and returns a numpy array of shape ``(N,)``. Ignored if ``label_path`` is already a numpy array. Defaults to ``None``. - log (bool): Apply a log transform to images before normalization. - Defaults to ``False``. - norm (str): Normalize images via "min-max" scaling ``[-1, 1]``, - or "center-scale". - two_dim (bool): If ``True``, reshape the data to treat each z-slice - as an independent 2D image. If ``False``, treat each sample as a - 3D volume. Defaults to ``True``. + reshape (str): If '2d' (default), treat each z-slice as independent 2D image + and reshape to (batch, channel, H, W). If '3d', reshape to + (batch, channel, D, H, W). If None, don't reshape. zthin (int): Thinning factor along the z axis, applied before - reshaping when ``two_dim=True``. A value of ``1`` applies no + reshaping when ``reshape='2d'``. A value of ``1`` applies no thinning. Defaults to ``1``. + normalization (str): Normalize image pixel distribution. + ['min-max', 'center-max'] + norm_kwargs: kwargs for pixel normalization function + transform (list): list of data transforms e.g. ['log', 'fft2']. + Ordered by operation. Note that 'log' always happens first if included. + read_only (bool): if True, only read the data from disk + to numpy array and return (no transform or normalization) Returns: - tuple: ``(images, labels)`` where ``images`` is a ``torch.Tensor`` on - ``device``, and ``labels`` is a LongTensor of shape ``(N,)`` or - ``None`` if ``label_path`` was not provided. + dict: A dictionary with keys + + * ``'images'`` — ``torch.Tensor`` on ``device``. + * ``'labels'`` — ``LongTensor`` of shape ``(N,)``, or ``None`` if + ``label_path`` was not provided. + * ``'norm'`` — fitted :class:`Normalization` instance, or ``None`` + if ``normalization`` was not requested. + * ``'tform'`` — fitted :class:`Transform` instance, or ``None`` if + ``transform`` was not requested. Example:: @@ -130,71 +136,246 @@ def load_data( def read_images(path): return np.load(path) - images, labels = load_data( + out = load_data( img_path="images.npy", img_read_fn=read_images, device="cpu", dtype=torch.float32, - two_dim=True, + reshape='2d', zthin=4, ) + images, labels = out['images'], out['labels'] """ - # --- images --------------------------------------------------------- - if isinstance(img_path, np.ndarray): - images = img_path - else: - images = img_read_fn(img_path) + # first check if feeding lists of filepaths + if isinstance(img_path, (list, tuple)): - if n_samples is not None: - if seed is None: - idx = slice(n_samples) - else: - rng = np.random.default_rng(seed) - idx = rng.choice(len(images), size=n_samples, replace=False) - images = images[idx] - else: - # images may be a memory-mapped array, so slice to instantiate it - images = images[:] + # list of filepaths: assumes label_path is also list if provided + n = len(img_path) + if label_path is None: + label_path = [label_path] * n + assert isinstance(label_path, (list, tuple)) + img_read_fn = img_read_fn if isinstance(img_read_fn, (list, tuple)) else [img_read_fn] * n + label_read_fn = label_read_fn if isinstance(label_read_fn, (list, tuple)) else [label_read_fn] * n + zthin = zthin if isinstance(zthin, (list, tuple)) else [zthin] * n + n_samples = n_samples if isinstance(n_samples, (list, tuple)) else [n_samples] * n + seed = seed if isinstance(seed, (list, tuple)) else [seed] * n + + # two options: norm per filepath, or one norm for all data + multi_norm = isinstance(normalization, (list, tuple)) + + if not multi_norm: + # one norm for all data. + # first load and concat data, then apply transform / norm + images, labels = [], [] + for img, ifn, lbl, lfn, zth, nsm, see in zip( + img_path, + img_read_fn, + label_path, + label_read_fn, + zthin, + n_samples, + seed, + ): + # only load and reshape the data + out = load_data( + img, + img_read_fn=ifn, + label_path=lbl, + label_read_fn=lfn, + device=device, + dtype=dtype, + reshape=reshape, + zthin=zth, + n_samples=nsm, + seed=see, + read_only=True, + ) + images.append(out['images']) + if out['labels'] is not None: + labels.append(out['labels']) + + # concat + images = torch.cat(images, dim=0) + if len(labels) > 0: + unq_labels = torch.cat([lbls[:1] for lbls in labels]) + labels = torch.cat(labels, dim=0) + else: + labels = None + + # pass through again to do transform / normalization + output = load_data( + images, + img_read_fn=None, # images already loaded + label_path=labels, + device=device, + dtype=dtype, + reshape=None, + normalization=normalization, + norm_kwargs=norm_kwargs, + transform=transform, + ) - images = torch.as_tensor(images, device=device, dtype=dtype) + return output - # --- labels --------------------------------------------------------- - labels = None - if label_path is not None: - if isinstance(label_path, np.ndarray): - labels = label_path else: - if label_read_fn is None: - raise ValueError( - "label_read_fn must be provided when label_path is a filepath." + # one norm for each dataset. + # load each data, apply tform / norm, then concat. + # note that this requires labels be supplied. + + # modify transform to be list of list if needed + if transform is not None: + # transform is either ['log', 'fft2'] or [['log'], None, ['log', 'fft2']] + if isinstance(transform[0], str): + transform = [transform] * n + + images, labels, norms, tforms = [], [], [], [] + for img, ifn, lbl, lfn, zth, nsm, see, nrm, nkw, trn in zip( + img_path, + img_read_fn, + label_path, + label_read_fn, + zthin, + n_samples, + seed, + normalization, + norm_kwargs, + transform + ): + out = load_data( + img, + img_read_fn=ifn, + label_path=lbl, + label_read_fn=lfn, + device=device, + dtype=dtype, + reshape=reshape, + zthin=zth, + n_samples=nsm, + seed=see, + normalization=nrm, + norm_kwargs=nkw, + transform=trn, + read_only=False, ) - labels = label_read_fn(label_path) + images.append(out['images']) + if out['labels'] is not None: + labels.append(out['labels']) + if out['tform'] is not None: + tforms.append(out['tform']) + if out['norm'] is not None: + norms.append(out['norm']) + + # concat + images = torch.cat(images, dim=0) + if len(labels) > 0: + unq_labels = torch.cat([lbls[0] for lbls in labels]) + labels = torch.cat(labels, dim=0) + else: + labels = None + + if len(tforms) > 0: + assert len(labels) > 0 + tform = MultiTransform(unq_labels, tforms) + else: + tform = None + + if len(norms) > 0: + assert len(labels) > 0 + norm = MultiNormalization(unq_labels, norms) + else: + norm = None + + output = { + 'images': images, + 'labels': labels, + 'norm': norm, + 'tform': tform, + } + + return output - if n_samples is not None: - labels = labels[idx] + else: + # assume img_path and label_path are single paths + # --- images --------------------------------------------------------- + if isinstance(img_path, (np.ndarray, torch.Tensor)): + images = img_path else: - labels = labels[:] + assert img_read_fn is not None + img_read_fn = globals()[img_read_fn] if isinstance(img_read_fn, str) else img_read_fn + images = img_read_fn(img_path) - labels = torch.as_tensor(labels, dtype=torch.long) + if n_samples is not None: + if seed is None: + idx = slice(n_samples) + else: + rng = np.random.default_rng(seed) + idx = rng.choice(len(images), size=n_samples, replace=False) + images = images[idx] + else: + if isinstance(images, np.ndarray) and not images.flags.writeable: + # may be memory-mapped ndarray, so make a copy + images = np.asarray(images, copy=True) + + images = torch.as_tensor(images, device=device, dtype=dtype) + + # --- labels --------------------------------------------------------- + labels = None + if label_path is not None: + if isinstance(label_path, (np.ndarray, torch.Tensor)): + labels = label_path + else: + assert label_read_fn is not None + label_read_fn = globals()[label_read_fn] if isinstance(label_read_fn, str) else label_read_fn + labels = label_read_fn(label_path) + + if n_samples is not None: + labels = labels[idx] + else: + labels = labels[:] + + labels = torch.as_tensor(labels, device=device, dtype=torch.long) + + # --- reshape -------------------------------------------------------- + if reshape == '2d': + images = images[:, ::zthin] + n_slices_per_vol = images.shape[1] + images = images.reshape(-1, 1, *images.shape[-2:]) + # broadcast labels per-slice so they remain 1:1 with the reshaped images + if labels is not None: + labels = labels.repeat_interleave(n_slices_per_vol, dim=0) + elif reshape == '3d': + images = images.unsqueeze(1) + + if read_only: + return {'images': images, 'labels': labels} # --- transforms ----------------------------------------------------- - if log: - images = images.log() - - if norm is not None: - if norm == 'center-scale': - images = center_scale_norm(images) - elif norm == 'min-max': - images = minmax_norm(images) - - # --- reshape -------------------------------------------------------- - if two_dim: - images = images[:, ::zthin] - images = images.reshape(-1, 1, *images.shape[-2:]) - else: - images = images.unsqueeze(1) + tform = None + if transform is not None: + from cosmodiff.transform import Transform + if 'log' in transform: + log = True + transform = [t for t in transform if t != 'log'] + else: + log = False + tform = Transform(transform, log=log) - return images, labels + images = tform(images) + + norm = None + if normalization is not None: + norm = Normalization(normalization, inplace=False, **(norm_kwargs or {})) + images = norm(images) + + + output = { + 'images': images, + 'labels': labels, + 'norm': norm, + 'tform': tform, + } + + return output def _import_class(qualified_name: str): @@ -272,45 +453,6 @@ def load_checkpoint(ckpt_path: str): return model, noise_scheduler, optimizer, lr_scheduler, augmentations -def minmax_norm(x: torch.Tensor) -> torch.Tensor: - """Normalize a tensor to ``[-1, 1]`` via min-max scaling. - - Args: - x (torch.Tensor): Input tensor of any shape. - - Returns: - torch.Tensor: Normalized tensor with values in ``[-1, 1]``. - """ - x = x - x.min() - x = x / x.max() - return x * 2 - 1 - - -def center_scale_norm(x: torch.Tensor, inplace: bool = False): - """Center a tensor based on its mean, and normalize by its absolute deviation. - - Args: - x (torch.Tensor): Input tensor of any shape. - inplace (bool): If True, edit inplace. - - Returns: - torch.Tensor: scaled tensor - float: avg - float: std - """ - # center - avg = x.mean() - if inplace: - x -= avg - else: - x = x - avg - - # scale by max-abs - x /= x.abs().max() - - return x - - def parse_config_model(config: dict): """Instantiate model, optimizer, noise_scheduler, and lr_scheduler from a parsed yaml config dict. Any missing keys return ``None``, which will @@ -320,15 +462,21 @@ def parse_config_model(config: dict): config (dict): Parsed yaml config, e.g. from ``yaml.safe_load()``. Returns: - tuple: ``(model, optimizer, noise_scheduler, lr_scheduler)`` where any - missing component is ``None``. + dict: A dictionary with keys ``'model'``, ``'optimizer'``, + ``'noise_scheduler'``, ``'lr_scheduler'``. Each entry is the + corresponding instantiated object, or ``None`` if the config did + not specify it (or its dependencies could not be built). Example:: with open("config.yaml") as f: config = yaml.safe_load(f) - model, optimizer, noise_scheduler, lr_scheduler = parse_model(config) + out = parse_config_model(config) + model = out['model'] + optimizer = out['optimizer'] + noise_scheduler = out['noise_scheduler'] + lr_scheduler = out['lr_scheduler'] """ import torch.optim.lr_scheduler as lr_schedulers @@ -360,29 +508,43 @@ def parse_config_model(config: dict): lr_cls = getattr(lr_schedulers, config["lr_scheduler"]["class"]) lr_scheduler = lr_cls(optimizer, **config["lr_scheduler"].get("kwargs", {})) - return model, optimizer, noise_scheduler, lr_scheduler + output = { + 'model': model, + 'optimizer': optimizer, + 'noise_scheduler': noise_scheduler, + 'lr_scheduler': lr_scheduler + } + + return output def parse_config_data(config: dict): """Load data and build an ``ArrayDataset`` from a parsed yaml config dict. - Read functions are resolved by name from ``cosmodiff.utils``. Any callable - in that module can be referenced by name in the yaml config. - Args: config (dict): Parsed yaml config, e.g. from ``yaml.safe_load()``. Returns: - ArrayDataset: Dataset ready to be passed to ``train()``. + dict: A dictionary with keys + + * ``'data'`` — :class:`ArrayDataset` ready to be passed to + ``train()``. + * ``'norm'`` — fitted :class:`Normalization` instance, or + ``None`` if normalization was not requested. + * ``'tform'`` — fitted :class:`Transform` instance, or ``None`` + if no ``transform`` list was given in the config. Example:: with open("config.yaml") as f: config = yaml.safe_load(f) - dataset = parse_config_data(config) + out = parse_config_data(config) + dataset = out['data'] + norm = out['norm'] + tform = out['tform'] """ - import cosmodiff.utils as utils_module + from cosmodiff import utils from cosmodiff.utils import ArrayDataset, load_data data_cfg = config["data"] @@ -392,66 +554,47 @@ def parse_config_data(config: dict): device = "cpu" if data_cfg.get("keep_on_cpu", False) \ else global_cfg.get("device", "cpu") - img_read_fn = getattr(utils_module, data_cfg["img_read_fn"]) - - label_path = data_cfg.get("label_path", None) - label_read_fn = None - if "label_read_fn" in data_cfg: - label_read_fn = getattr(utils_module, data_cfg["label_read_fn"]) - img_path = data_cfg["img_path"] - n_samples = data_cfg.get("n_samples", None) - seed = data_cfg.get("seed", None) - log = data_cfg.get("log", False) - - _load_kwargs = dict( - img_read_fn=img_read_fn, + out = load_data( + img_path=data_cfg["img_path"], + img_read_fn=data_cfg.get('img_read_fn', None), device=device, dtype=dtype, - label_read_fn=label_read_fn, - norm=data_cfg.get("norm", 'center-scale'), - two_dim=data_cfg.get("two_dim", True), + label_path=data_cfg.get("label_path", None), + label_read_fn=data_cfg.get('label_read_fn', None), + reshape=data_cfg.get("reshape", '2d'), zthin=data_cfg.get("zthin", 1), + n_samples=data_cfg.get("n_samples", None), + seed=data_cfg.get("seed", None), + normalization=data_cfg.get('normalization', None), + norm_kwargs=data_cfg.get('norm_kwargs', None), + transform=data_cfg.get('transform', None), ) - if isinstance(img_path, (list, tuple)): - n = len(img_path) - label_paths = label_path if isinstance(label_path, (list, tuple)) else [label_path] * n - n_samples_list = n_samples if isinstance(n_samples, (list, tuple)) else [n_samples] * n - seeds_list = seed if isinstance(seed, (list, tuple)) else [seed] * n - log_list = log if isinstance(log, (list, tuple)) else [log] * n - - all_images, all_labels = [], [] - for p, lp, ns, sd, lg in zip(img_path, label_paths, n_samples_list, seeds_list, log_list): - imgs, lbls = load_data(img_path=p, label_path=lp, n_samples=ns, seed=sd, log=lg, **_load_kwargs) - all_images.append(imgs) - if lbls is not None: - all_labels.append(lbls) - - images = torch.cat(all_images, dim=0) - labels = torch.cat(all_labels, dim=0) if all_labels else None - else: - images, labels = load_data( - img_path=img_path, - label_path=label_path, - n_samples=n_samples, - seed=seed, - log=log, - **_load_kwargs, - ) - augmentations = None if "augmentations" in config: from cosmodiff.augment import config_augmentations augmentations = config_augmentations(config["augmentations"]) - return ArrayDataset(images, labels=labels, augmentations=augmentations) + output = { + 'data': ArrayDataset( + out['images'], labels=out['labels'], augmentations=augmentations, + ), + 'norm': out['norm'], + 'tform': out['tform'], + } + + return output def npy_read_fn(fname): return np.load(fname, mmap_mode='r') +def txt_read_fn(fname): + return np.loadtxt(fname) + + def read_logs(output_dir: str) -> dict: """Extract training metrics from TensorBoard logs produced by ``optim.train()``. @@ -537,8 +680,8 @@ def plot_metrics(metrics: dict | str, save_dir: str = None, show: bool = False) Example:: # from dict - metrics = train(dataset, model) - plot_metrics(metrics, show=True) + result = train(dataset, model) + plot_metrics(result['metrics'], show=True) # from file plot_metrics("output/metrics_epoch_49.json", save_dir="output/plots") @@ -609,6 +752,42 @@ def plot_metrics(metrics: dict | str, save_dir: str = None, show: bool = False) return fig1, fig2, fig3 +def plot_ema_profiles(profiles: dict, ax=None): + """Plot per-checkpoint and synthesized EMA weight profiles. + + Takes the dict returned by ``cosmodiff.optim.compute_ema_profiles`` and + produces a figure analogous to Fig. 4 of Karras et al. 2024. + + Args: + profiles: output of ``compute_ema_profiles``. + ax: existing Axes to draw on; creates a new figure if ``None``. + + Returns: + matplotlib Axes. + """ + import matplotlib.pyplot as plt + + if ax is None: + fig, ax = plt.subplots() + + epochs = profiles['epochs'] + + for t_i, s_j, w in profiles['basis']: + ax.plot(epochs, w, color='steelblue', alpha=0.4, linewidth=0.8, + label=rf'checkpoint epoch {t_i}, $\sigma_{{rel}}$={s_j}') + + ax.plot(epochs, profiles['target'], 'k--', linewidth=1.5, + label=r'target profile') + ax.plot(epochs, profiles['synthesized'], 'r-', linewidth=2, + label=r'synthesized profile') + + ax.set_xlabel('Training epoch') + ax.set_ylabel('EMA weight (normalized)') + ax.legend(fontsize=7) + + return ax + + def find_latest_checkpoint(output_dir: str) -> str | None: """Return the path to the latest checkpoint in ``output_dir``, or ``None`` if no checkpoints exist. diff --git a/cosmodiff/version.py b/cosmodiff/version.py deleted file mode 100644 index 9f5f846..0000000 --- a/cosmodiff/version.py +++ /dev/null @@ -1,4 +0,0 @@ -__version__ = "0.0.1" - -def version(): - return "cosmodiff v{}".format(__version__) diff --git a/docs/_static/banner.png b/docs/_static/banner.png new file mode 100644 index 0000000..39e3698 Binary files /dev/null and b/docs/_static/banner.png differ diff --git a/docs/_static/gen_cosmodiff_banner.py b/docs/_static/gen_cosmodiff_banner.py new file mode 100644 index 0000000..72c7db8 --- /dev/null +++ b/docs/_static/gen_cosmodiff_banner.py @@ -0,0 +1,130 @@ +"""Generate a realistic cosmic-web banner for cosmo_diffusion. + +Left half: cosmic web realization A. +Right half: a different realization B. +Middle: Gaussian-enveloped noise that smoothly takes over near the centre, +emphasising the diffusion process A → noise → B. +""" + +import numpy as np +import matplotlib.pyplot as plt +import cmasher as cmr + + +def power_law_field(shape, alpha=4.5, seed=0): + """Gaussian random field with power spectrum P(k) ~ k^{-alpha}.""" + h, w = shape + rng = np.random.default_rng(seed) + + noise = rng.standard_normal((h, w)) + F = np.fft.fft2(noise) + + kx = np.fft.fftfreq(w) + ky = np.fft.fftfreq(h) + KX, KY = np.meshgrid(kx, ky) + K = np.sqrt(KX ** 2 + KY ** 2) + K[0, 0] = 1.0 # avoid div-by-zero + + P = K ** (-alpha) + P[0, 0] = 0.0 + F_shaped = F * np.sqrt(P) + field = np.real(np.fft.ifft2(F_shaped)) + return field + + +def cosmic_web_field(shape, seed): + """A cosmic-web-ish 2D field: diffuse filamentary background + bright halos.""" + rng = np.random.default_rng(seed) + h, w = shape + + # Filamentary diffuse component + field = power_law_field(shape, alpha=2.2, seed=seed) + field = (field - field.mean()) / field.std() + field = np.exp(0.75 * field) - 1.0 # log-normal stretch — emphasises peaks + field = np.clip(field, 0, None) + + # Add bright halos + #Y, X = np.mgrid[:h, :w] + #n_halos = int(rng.integers(20, 32)) + #for _ in range(n_halos): + # cx = rng.uniform(0, w) + # cy = rng.uniform(0, h) + # sigma = rng.uniform(3.0, 14.0) + # amp = rng.uniform(0.8, 4.0) + # field += amp * np.exp(-((X - cx) ** 2 + (Y - cy) ** 2) / (2 * sigma ** 2)) + + # Normalise to [0, 1] (saturate the brightest tail for punch). Using a + # slightly lower percentile boosts midrange brightness so the diffuse + # filamentary structure remains visible. + field = field / np.percentile(field, 95.0) + return np.clip(field, 0.0, 1.0) + + +def make_banner(out_path, W=1200, H=300): + print(f"generating {W} x {H} banner …") + + # Two distinct cosmic-web realisations (seeds chosen so the two halves + # have comparable bright-feature density). + left = cosmic_web_field((H, W), seed=5) + right = cosmic_web_field((H, W), seed=4) + + # Noise with cosmic-web-like dynamic range (so the blend looks coherent) + rng = np.random.default_rng(999) + noise = rng.standard_normal((H, W)) + noise = (noise - noise.min()) / (noise.max() - noise.min()) + noise = noise ** 1.4 + noise = noise * 1 # slightly dimmer than the bright cosmic peaks + + # Spatial weights along the x-axis + x = np.arange(W) + # Gaussian noise envelope — peaks in the middle, falls off toward both edges + sigma_n = W / 5.0 + noise_env = np.exp(-((x - W / 2) ** 2) / (2 * sigma_n ** 2)) + + # Left → right taper using a smooth sigmoid + transition = W / 9.0 + left_taper = 1.0 / (1.0 + np.exp((x - W / 2) / transition)) + right_taper = 1.0 - left_taper + + left_w = (1.0 - noise_env) * left_taper + right_w = (1.0 - noise_env) * right_taper + noise_w = noise_env + + # Compose + composite = left * left_w[None, :] + right * right_w[None, :] + noise * noise_w[None, :] + + # CMasher's colormap + cmap = cmr.freeze + + fig, ax = plt.subplots(figsize=(W / 150, H / 150), dpi=150) + ax.imshow( + composite, + cmap=cmap, + vmin=0.0, + vmax=1.0, + aspect="auto", + interpolation="bicubic", + ) + # Title at the bottom-center + ax.text( + W / 2, + H - H * 0.06, + "cosmo_diffusion", + color="#ffffff", + fontsize=34, + fontname="Courier New", + weight="bold", + ha="center", + va="bottom", + alpha=0.95, + ) + ax.axis("off") + fig.subplots_adjust(0, 0, 1, 1) + fig.savefig(out_path, bbox_inches="tight", pad_inches=0) + plt.close(fig) + + print(f"wrote {out_path}") + + +if __name__ == "__main__": + make_banner("/Users/nkern/Software/cosmo_diffusion/banner.png") diff --git a/pyproject.toml b/pyproject.toml index f7fde62..14e9c08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ classifiers = [ "Programming Language :: Python :: 3", "Intended Audience :: Science/Research", ] -dynamic = ["version"] +version = "0.0.1" [project.optional-dependencies] dev = [ @@ -29,6 +29,7 @@ dev = [ "pytest", "pytest-cov", "coverage", + "ema-pytorch", ] docs = [ "cosmodiff[dev]", @@ -48,5 +49,12 @@ testpaths = 'cosmodiff/tests' filterwarnings = ["ignore::DeprecationWarning"] [tool.setuptools] -script-files = ["scripts/cosmodiff_train.py"] +script-files = ["scripts/cosmodiff_train.py", "scripts/cosmodiff_sample.py"] + +[tool.setuptools.packages.find] +include = ["cosmodiff*"] +exclude = ["cosmodiff.tests*"] + +[tool.setuptools.package-data] +"cosmodiff.data" = ["*.npy", "*.yaml", "*.txt"] diff --git a/scripts/cosmodiff_sample.py b/scripts/cosmodiff_sample.py index 43a8668..cf6ab91 100644 --- a/scripts/cosmodiff_sample.py +++ b/scripts/cosmodiff_sample.py @@ -4,6 +4,15 @@ Usage: python cosmodiff_sample.py --checkpoint path/to/checkpoint --n_samples 100 --output samples.npy python cosmodiff_sample.py --output_dir path/to/run --n_samples 64 --image_shape 1 64 64 + +Fast sampling (diffusion-trained models): + --scheduler DPMSolverMultistepScheduler --num_steps 25 + +Fast sampling (flow-matching-trained models): + --scheduler FlowMatchHeunDiscreteScheduler --num_steps 25 + +EMA: + --ema_sigma_rel 0.05 → synthesize EMA at target sigma_rel (uses all checkpoints in output_dir) """ import argparse @@ -56,15 +65,81 @@ def main(): type=int, nargs="+", default=None, - help="Class labels for conditional generation. Must be length 1 " - "(broadcast to all samples) or n_samples.", + help="Discrete class labels for conditional generation. Must be length 1 " + "(broadcast to all samples) or n_samples. Use --continuous_labels for " + "encoder_hidden_states-style conditioning.", + ) + parser.add_argument( + "--continuous_labels", + type=str, + default=None, + help="Path to a .npy file of float conditioning vectors of shape " + "(n_samples, D) or (1, D) (broadcast). Used with --conditioning continuous.", + ) + parser.add_argument( + "--conditioning", + type=str, + choices=["discrete", "continuous"], + default="discrete", + help="Conditioning mode used during training. Defaults to 'discrete'.", + ) + parser.add_argument( + "--guidance_scale", + type=float, + default=None, + help="Classifier-free guidance scale at inference. None (default) disables " + "CFG amplification. Typical values 1.0-7.0.", + ) + parser.add_argument( + "--ema_sigma_rel", + type=float, + default=None, + help="Target sigma_rel for post-hoc EMA weight synthesis. When set, the " + "model weights are replaced by the synthesized EMA (across all " + "checkpoints in output_dir) before sampling.", + ) + parser.add_argument( + "--scheduler", + type=str, + default=None, + help="Name of a diffusers scheduler class to use at inference, e.g. " + "'DDIMScheduler', 'HeunDiscreteScheduler', 'DPMSolverMultistepScheduler'. " + "Defaults to the training scheduler. The new scheduler is built via " + ".from_config() to inherit the trained beta schedule and prediction_type.", ) parser.add_argument( - "--ddim_thinning", + "--num_steps", type=int, default=None, - help="Thinning factor to reduce inference steps (e.g. 10 → 100 steps from 1000). " - "Automatically switches to DDIMScheduler regardless of what was used during training.", + help="Number of inference steps. Defaults to the scheduler's " + "num_train_timesteps (full schedule).", + ) + parser.add_argument( + "--s_churn", + type=float, + default=None, + help="EDM-style stochasticity injection (Karras et al. 2022). " + "0 = pure ODE; larger = more SDE-like. Only consumed by " + "Euler/Heun-family schedulers (diffusion or FM); silently " + "ignored for others (DDPM, DDIM, DPM-Solver, etc.).", + ) + parser.add_argument( + "--s_tmin", + type=float, + default=None, + help="Lower-bound timestep for churn gating; ignored when --s_churn is unset.", + ) + parser.add_argument( + "--s_tmax", + type=float, + default=None, + help="Upper-bound timestep for churn gating; ignored when --s_churn is unset.", + ) + parser.add_argument( + "--s_noise", + type=float, + default=None, + help="Multiplier on injected noise magnitude during churn (default 1.0).", ) parser.add_argument( "--seed", @@ -87,27 +162,30 @@ def main(): args = parser.parse_args() from cosmodiff.utils import load_checkpoint, find_latest_checkpoint - from cosmodiff.optim import generate + from cosmodiff.optim import generate, synthesize_ema_from_checkpoints # --- resolve checkpoint --------------------------------------------- if args.checkpoint is not None: ckpt_path = args.checkpoint + output_dir = os.path.dirname(os.path.abspath(ckpt_path)) else: - ckpt_path = find_latest_checkpoint(args.output_dir) + output_dir = args.output_dir + ckpt_path = find_latest_checkpoint(output_dir) if ckpt_path is None: - raise FileNotFoundError(f"No checkpoints found in {args.output_dir}") + raise FileNotFoundError(f"No checkpoints found in {output_dir}") if args.verbose: print(f"Loading checkpoint: {ckpt_path}") model, noise_scheduler, _optimizer, _lr_scheduler, _augmentations = load_checkpoint(ckpt_path) - # --- swap to DDIM if thinning requested ----------------------------- - if args.ddim_thinning is not None: - from diffusers import DDIMScheduler - noise_scheduler = DDIMScheduler.from_config(noise_scheduler.config) + # --- swap to a different scheduler if requested --------------------- + if args.scheduler is not None: + import diffusers + sched_cls = getattr(diffusers, args.scheduler) + noise_scheduler = sched_cls.from_config(noise_scheduler.config) if args.verbose: - print("Switched to DDIMScheduler for fast sampling.") + print(f"Switched to {args.scheduler} for inference.") # --- device --------------------------------------------------------- if args.device is not None: @@ -121,6 +199,16 @@ def main(): if args.verbose: print(f"Using device: {device}") + # --- EMA synthesis (overrides model weights) ------------------------ + if args.ema_sigma_rel is not None: + if args.verbose: + print(f"Synthesizing EMA weights at sigma_rel={args.ema_sigma_rel} from {output_dir}") + model = synthesize_ema_from_checkpoints( + model, output_dir, sigma_rel_target=args.ema_sigma_rel, + ) + model = model.to(device) + model.eval() + # --- image shape ---------------------------------------------------- if args.image_shape is not None: image_shape = tuple(args.image_shape) @@ -138,16 +226,35 @@ def main(): # --- labels --------------------------------------------------------- labels_tensor = None - if args.labels is not None: - if len(args.labels) == 1: - labels_tensor = torch.full((args.n_samples,), args.labels[0], dtype=torch.long) - elif len(args.labels) == args.n_samples: - labels_tensor = torch.tensor(args.labels, dtype=torch.long) - else: + if args.conditioning == "discrete": + if args.labels is not None: + if len(args.labels) == 1: + labels_tensor = torch.full((args.n_samples,), args.labels[0], dtype=torch.long) + elif len(args.labels) == args.n_samples: + labels_tensor = torch.tensor(args.labels, dtype=torch.long) + else: + raise ValueError( + f"--labels must be length 1 or n_samples ({args.n_samples}), " + f"got {len(args.labels)}." + ) + else: # continuous + if args.continuous_labels is None: + raise ValueError( + "--continuous_labels is required when --conditioning continuous." + ) + cont_arr = np.load(args.continuous_labels) + if cont_arr.ndim != 2: + raise ValueError( + f"--continuous_labels must be 2D (N, D); got shape {cont_arr.shape}." + ) + if cont_arr.shape[0] == 1: + cont_arr = np.broadcast_to(cont_arr, (args.n_samples, cont_arr.shape[1])) + elif cont_arr.shape[0] != args.n_samples: raise ValueError( - f"--labels must be length 1 or n_samples ({args.n_samples}), " - f"got {len(args.labels)}." + f"--continuous_labels first dim must be 1 or n_samples ({args.n_samples}), " + f"got {cont_arr.shape[0]}." ) + labels_tensor = torch.as_tensor(np.asarray(cont_arr), dtype=torch.float32) # --- generator ------------------------------------------------------ generator = None @@ -172,7 +279,13 @@ def main(): batch_size=bs, image_shape=image_shape, labels=batch_labels, - ddim_thinning=args.ddim_thinning, + guidance_scale=args.guidance_scale, + conditioning=args.conditioning, + num_steps=args.num_steps, + s_churn=args.s_churn, + s_tmin=args.s_tmin, + s_tmax=args.s_tmax, + s_noise=args.s_noise, device=device, generator=generator, ) @@ -184,7 +297,9 @@ def main(): result = np.concatenate(all_samples, axis=0) - os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + out_dir = os.path.dirname(os.path.abspath(args.output)) + if out_dir: + os.makedirs(out_dir, exist_ok=True) np.save(args.output, result) print(f"Saved {result.shape} array to {args.output}") diff --git a/scripts/cosmodiff_train.py b/scripts/cosmodiff_train.py index 8467239..b17df94 100644 --- a/scripts/cosmodiff_train.py +++ b/scripts/cosmodiff_train.py @@ -75,37 +75,39 @@ def main(): # --- check for existing checkpoint ---------------------------------- latest_ckpt = find_latest_checkpoint(output_dir) - dataset = parse_config_data(config) + data_out = parse_config_data(config) + dataset = data_out["data"] if latest_ckpt is not None: print(f"Resuming from checkpoint: {latest_ckpt}") - metrics = train( + result = train( dataset, resume_from_checkpoint=latest_ckpt, output_dir=output_dir, **config["train"], ) + else: print("No checkpoint found, training from scratch.") - model, optimizer, noise_scheduler, lr_scheduler = parse_config_model(config) - metrics = train( + model_out = parse_config_model(config) + result = train( dataset, - model, - optimizer=optimizer, - noise_scheduler=noise_scheduler, - lr_scheduler=lr_scheduler, + model_out["model"], + optimizer=model_out["optimizer"], + noise_scheduler=model_out["noise_scheduler"], + lr_scheduler=model_out["lr_scheduler"], output_dir=output_dir, **config["train"], ) print(f"Training complete.") - print(f"Final epoch loss: {metrics['epoch_loss'][-1]:.4f}") - print(f"Total time: {sum(metrics['epoch_times']):.1f}s") + print(f"Final epoch loss: {result['metrics']['epoch_loss'][-1]:.4f}") + print(f"Total time: {sum(result['metrics']['epoch_times']):.1f}s") metrics_path = os.path.join( - output_dir, "metrics_epoch_{:04d}.json".format(len(metrics["epoch_loss"]) - 1) + output_dir, "metrics_epoch_{:04d}.json".format(len(result['metrics']["epoch_loss"]) - 1) ) - write_metrics(metrics, metrics_path) + write_metrics(result['metrics'], metrics_path) print(f"Metrics written to {metrics_path}") if __name__ == "__main__": diff --git a/setup.py b/setup.py deleted file mode 100644 index 85b1725..0000000 --- a/setup.py +++ /dev/null @@ -1,35 +0,0 @@ -from setuptools import setup -import os -import ast - -# get version from __init__.py -init_file = os.path.join('/'.join(os.path.abspath(__file__).split('/')[:-1]), 'cosmodiff/version.py') -with open(init_file, 'r') as f: - lines = f.readlines() - for l in lines: - if "__version__ =" in l: - version = ast.literal_eval(l.split('=')[1].strip()) - -def package_files(package_dir, subdirectory): - # walk the input package_dir/subdirectory - # return a package_data list - paths = [] - directory = os.path.join(package_dir, subdirectory) - for (path, directories, filenames) in os.walk(directory): - for filename in filenames: - path = path.replace(package_dir + '/', '') - paths.append(os.path.join(path, filename)) - return paths - - -data_files = package_files('cosmodiff', 'data') + package_files('cosmodiff', 'config') - -setup( - version = version, - license = 'MIT', - package_data = {'cosmodiff': data_files}, - include_package_data = True, - packages = ['cosmodiff'], - package_dir = {'cosmo_diffusion': 'cosmodiff'}, - zip_safe = False - )