Index

A

Activation (class in mindscience.e3nn.nn) AdaHessian (class in mindscience.common) add_noise() (mindscience.diffuser.DiffusionScheduler method) all_to_all_from_hidden_to_sequence() (in module mindscience.distributed.mappings) all_to_all_from_sequence_to_hidden() (in module mindscience.distributed.mappings) angles_to_matrix() (in module mindscience.e3nn.o3) angles_to_xyz() (in module mindscience.e3nn.o3) ASD_FFT (class in mindscience.sciops.fft) asd_fft() (in module mindscience.sciops.fft) ASD_FFT2D (class in mindscience.sciops.fft) asd_fft2d() (in module mindscience.sciops.fft) asd_fftn() (in module mindscience.sciops.fft) ASD_IFFT (class in mindscience.sciops.fft) asd_ifft() (in module mindscience.sciops.fft)

ASD_IFFT2D (class in mindscience.sciops.fft) asd_ifft2d() (in module mindscience.sciops.fft) asd_ifftn() (in module mindscience.sciops.fft) ASD_IRFFT (class in mindscience.sciops.fft) asd_irfft() (in module mindscience.sciops.fft) ASD_IRFFT2D (class in mindscience.sciops.fft) asd_irfft2d() (in module mindscience.sciops.fft) asd_irfftn() (in module mindscience.sciops.fft) ASD_RFFT (class in mindscience.sciops.fft) asd_rfft() (in module mindscience.sciops.fft) ASD_RFFT2D (class in mindscience.sciops.fft) asd_rfft2d() (in module mindscience.sciops.fft) asd_rfftn() (in module mindscience.sciops.fft) Attention (class in mindscience.models.transformer)

B

batched_hessian() (in module mindscience.common)

batched_jacobian() (in module mindscience.common) BatchNorm (class in mindscience.e3nn.nn)

C

change_basis_real_to_complex() (in module mindscience.e3nn.o3) check_dict_type() (in module mindscience.utils) check_dict_type_value() (in module mindscience.utils) check_dict_value() (in module mindscience.utils) check_param_type() (in module mindscience.utils) check_param_type_value() (in module mindscience.utils) check_param_value() (in module mindscience.utils) ColumnParallelLinear (class in mindscience.distributed.modules) compose_angles() (in module mindscience.e3nn.o3)

ConditionDiffusionTransformer (class in mindscience.diffuser) Cone (class in mindscience.data) copy_to_all() (in module mindscience.distributed.mappings) count() (mindscience.e3nn.o3.Irreps method) CSGDifference (class in mindscience.data) CSGIntersection (class in mindscience.data) CSGUnion (class in mindscience.data) CSGXOR (class in mindscience.data) Cuboid (class in mindscience.data) Cylinder (class in mindscience.data)

D

DCT (class in mindscience.sciops.dft) DDIMPipeline (class in mindscience.diffuser) DDIMScheduler (class in mindscience.diffuser) DDPMPipeline (class in mindscience.diffuser) DDPMScheduler (class in mindscience.diffuser) decompose() (mindscience.e3nn.o3.Irreps method)

DFTn (class in mindscience.sciops.dft) DiffusionScheduler (class in mindscience.diffuser) DiffusionTrainer (class in mindscience.diffuser) DiffusionTransformer (class in mindscience.diffuser) Disk (class in mindscience.data) DropPath (class in mindscience.models.layers) DST (class in mindscience.sciops.dft)

E

Einsum (class in mindscience.sciops.einsum)

ElementwiseTensorProduct (class in mindscience.e3nn.o3) evo_attention() (in module mindscience.sciops.evoformer_attention)

F

FCSequential (class in mindscience.models.layers) FFNO (class in mindscience.models.neural_operator) FFNO1D (class in mindscience.models.neural_operator) FFNO2D (class in mindscience.models.neural_operator) FFNO3D (class in mindscience.models.neural_operator) FFNOBlocks (class in mindscience.models.neural_operator) filter() (mindscience.e3nn.o3.Irreps method)

FixedPoint (class in mindscience.data) FNO1D (class in mindscience.models.neural_operator) FNO2D (class in mindscience.models.neural_operator) FNO3D (class in mindscience.models.neural_operator) FNOBlocks (class in mindscience.models.neural_operator) FullTensorProduct (class in mindscience.e3nn.o3) FullyConnectedNet (class in mindscience.e3nn.nn) FullyConnectedTensorProduct (class in mindscience.e3nn.o3)

G

Gate (class in mindscience.e3nn.nn) gather_from_hidden() (in module mindscience.distributed.mappings) gather_from_sequence() (in module mindscience.distributed.mappings) generate_sampling_config() (in module mindscience.data) Geometry (class in mindscience.data) GeometryWithTime (class in mindscience.data) get_2d_sin_cos_pos_embed() (in module mindscience.common) get_activation() (in module mindscience.models.layers) get_context_parallel_group() (in module mindscience.distributed.manager) get_context_parallel_group_name() (in module mindscience.distributed.manager) get_context_parallel_rank() (in module mindscience.distributed.manager) get_context_parallel_world_size() (in module mindscience.distributed.manager) get_data_context_parallel_group() (in module mindscience.distributed.manager) get_data_context_parallel_group_name() (in module mindscience.distributed.manager) get_data_context_parallel_rank() (in module mindscience.distributed.manager)

get_data_context_parallel_world_size() (in module mindscience.distributed.manager) get_data_parallel_group() (in module mindscience.distributed.manager) get_data_parallel_group_name() (in module mindscience.distributed.manager) get_data_parallel_rank() (in module mindscience.distributed.manager) get_data_parallel_world_size() (in module mindscience.distributed.manager) get_loss() (mindscience.diffuser.DiffusionTrainer method) get_loss_metric() (in module mindscience.common) get_multi_step_lr() (in module mindscience.common) get_poly_lr() (in module mindscience.common) get_tensor_parallel_group() (in module mindscience.distributed.manager) get_tensor_parallel_group_name() (in module mindscience.distributed.manager) get_tensor_parallel_rank() (in module mindscience.distributed.manager) get_tensor_parallel_world_size() (in module mindscience.distributed.manager) get_warmup_cosine_annealing_lr() (in module mindscience.common) glorot_uniform() (in module mindscience.common) GraphCastNet (class in mindscience.models.GraphCast)

H

HyperCube (class in mindscience.data)

I

IDCT (class in mindscience.sciops.dft) identity_angles() (in module mindscience.e3nn.o3) IDFTn (class in mindscience.sciops.dft) IDST (class in mindscience.sciops.dft) init_edge_rot_mat() (in module mindscience.e3nn.so2_conv) initialize_affine_weight() (in module mindscience.distributed.modules)

initialize_parallel() (in module mindscience.distributed.manager) InputScale (class in mindscience.models.layers) Interval (class in mindscience.data) IRDFTn (class in mindscience.sciops.dft) Irrep (class in mindscience.e3nn.o3) Irreps (class in mindscience.e3nn.o3) is_scalar() (mindscience.e3nn.o3.Irrep method)

L

lecun_init() (in module mindscience.common) Linear (class in mindscience.e3nn.o3) LinearBias (class in mindscience.e3nn.o3)

LinearBlock (class in mindscience.models.layers) load_yaml_config() (in module mindscience.utils) log_config() (in module mindscience.utils) log_timer() (in module mindscience.utils)

M

MaskedLayerNorm (class in mindscience.models.layers) matrix_to_angles() (in module mindscience.e3nn.o3) matrix_x() (in module mindscience.e3nn.o3) matrix_y() (in module mindscience.e3nn.o3)

matrix_z() (in module mindscience.e3nn.o3) MTLWeightedLoss (class in mindscience.common) MultiHeadAttention (class in mindscience.models.transformer) MultiScaleFCSequential (class in mindscience.models.layers)

N

narrow() (mindscience.e3nn.so2_conv.SO3Rotation static method) Ncon (class in mindscience.e3nn.utils)

Norm (class in mindscience.e3nn.o3) NormActivation (class in mindscience.e3nn.nn)

O

OneHot (class in mindscience.e3nn.nn)

P

PartSamplingConfig (class in mindscience.data) patchify() (in module mindscience.common) Pentagon (class in mindscience.data) pixel_shuffle() (in module mindscience.common)

pixel_unshuffle() (in module mindscience.common) PixelShuffle (class in mindscience.common) PixelUnshuffle (class in mindscience.common) Polygon (class in mindscience.data) print_log() (in module mindscience.utils)

R

radius() (in module mindscience.e3nn.utils) radius_full() (in module mindscience.e3nn.utils) radius_graph() (in module mindscience.e3nn.utils) radius_graph_full() (in module mindscience.e3nn.utils) rand_angles() (in module mindscience.e3nn.o3) randn() (mindscience.e3nn.o3.Irreps method) RDFTn (class in mindscience.sciops.dft) Rectangle (class in mindscience.data)

reduce_from_all() (in module mindscience.distributed.mappings) reduce_scatter_to_sequence() (in module mindscience.distributed.mappings) RelativeRMSELoss (class in mindscience.common) remove_zero_multiplicities() (mindscience.e3nn.o3.Irreps method) ResBlock (class in mindscience.models.layers) rotate() (mindscience.e3nn.so2_conv.SO3Rotation method) rotate_inv() (mindscience.e3nn.so2_conv.SO3Rotation method) rotation_to_wigner_d_matrix() (mindscience.e3nn.so2_conv.SO3Rotation static method) RowParallelLinear (class in mindscience.distributed.modules)

S

sampling() (mindscience.data.Disk method) (mindscience.data.FixedPoint method) (mindscience.data.GeometryWithTime method) (mindscience.data.HyperCube method) SamplingConfig (class in mindscience.data) Scatter (class in mindscience.e3nn.nn) scatter_to_hidden() (in module mindscience.distributed.mappings) scatter_to_sequence() (in module mindscience.distributed.mappings) set_fft_cache_size() (in module mindscience.sciops.fft) set_name() (mindscience.data.Geometry method) set_sampling_config() (mindscience.data.Geometry method) (mindscience.data.GeometryWithTime method) set_timesteps() (mindscience.diffuser.DDPMScheduler method) (mindscience.diffuser.DiffusionScheduler method) set_wigner() (mindscience.e3nn.so2_conv.SO3Rotation method) simplify() (mindscience.e3nn.o3.Irreps method)

SO2Convolution (class in mindscience.e3nn.so2_conv) so3_generators() (in module mindscience.e3nn.o3) SO3Rotation (class in mindscience.e3nn.so2_conv) soft_one_hot_linspace() (in module mindscience.e3nn.nn) soft_unit_step() (in module mindscience.e3nn.nn) SoftOneHotLinspace (class in mindscience.e3nn.nn) sort() (mindscience.e3nn.o3.Irreps method) SpectralNorm (class in mindscience.common) spherical_harmonics() (in module mindscience.e3nn.o3) (mindscience.e3nn.o3.Irreps static method) SphericalHarmonics (class in mindscience.e3nn.o3) SReLU (class in mindscience.models.layers.activation) step() (mindscience.diffuser.DDIMScheduler method) (mindscience.diffuser.DDPMScheduler method) (mindscience.diffuser.DiffusionScheduler method) su2_generators() (in module mindscience.e3nn.o3)

T

TensorProduct (class in mindscience.e3nn.o3) TensorSquare (class in mindscience.e3nn.o3) Tetrahedron (class in mindscience.data) TimeDomain (class in mindscience.data)

to_2tuple() (in module mindscience.common) to_3tuple() (in module mindscience.common) TransformerBlock (class in mindscience.models.transformer) Triangle (class in mindscience.data)

U

UNet2D (class in mindscience.models.layers)

unpatchify() (in module mindscience.common)

V

VisionTransformer (class in mindscience.models.transformer)

W

WaveletTransformLoss (class in mindscience.common) wigD_from_angles() (mindscience.e3nn.o3.Irrep method) (mindscience.e3nn.o3.Irreps method)

wigD_from_matrix() (mindscience.e3nn.o3.Irrep method) (mindscience.e3nn.o3.Irreps method) wigner_3j() (in module mindscience.e3nn.o3) wigner_D() (in module mindscience.e3nn.o3)

X

xyz_to_angles() (in module mindscience.e3nn.o3)