Skip to content

Add square-symmetry augmentations and sampling fixes#8

Open
JiamingPan wants to merge 1 commit into
mainfrom
codex/augmentations
Open

Add square-symmetry augmentations and sampling fixes#8
JiamingPan wants to merge 1 commit into
mainfrom
codex/augmentations

Conversation

@JiamingPan
Copy link
Copy Markdown
Collaborator

Summary:

  • add RandomRot90 and RandomDihedral2D augmentations for square-symmetry physics augmentation
  • keep RandomFlip config dims as a tuple and add augmentation unit tests
  • make sampling compatible with schedulers whose step() does not accept generator, avoid duplicate EMA symlink failures, and unwrap synthesized EMA models before sampling

Notes:

  • This is intended for CAMELS-style scalar maps where rotations/flips are physical symmetries and should not change the field distribution.
  • I could not run full pytest locally because this machine lacks torch/pytest; please run the existing test suite in the normal cosmo_diffusion environment or CI before merging.

@JiamingPan JiamingPan requested a review from nkern May 17, 2026 18:17
@nkern nkern marked this pull request as ready for review May 18, 2026 02:14
@nkern
Copy link
Copy Markdown
Owner

nkern commented May 18, 2026

thanks for the addition jiaming--will have a look and give comments. fyi--the CI automatically runs pytest and can import torch, and it looks like it passes! but will leave some comments before we merge

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants