Revert "Rename package to ai2modulus"

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
           "--ignore-magic", "--fail-under=99", "--exclude=['setup.py', 'test', 'build', 'docs']",  
           "--ignore-regex=['forward', 'reset_parameters', 'extra_repr', 'MetaData', 'apply_activation','exec_activation']", 
           "--color", "--"]
-      exclude: ^ai2modulus/internal/
+      exclude: ^modulus/internal/
 
 -   repo: local
     hooks:

Dockerfile

@@ -28,16 +28,16 @@ ENV _CUDA_COMPAT_TIMEOUT=90
 # Install custom onnx
 # TODO: Find a fix to eliminate the custom build
 # Forcing numpy update to over ride numba 0.56.4 max numpy constraint
-COPY . /ai2modulus/ 
-RUN if [ -e "/ai2modulus/deps/onnxruntime_gpu-1.14.0-cp38-cp38-linux_x86_64.whl" ]; then \
+COPY . /modulus/ 
+RUN if [ -e "/modulus/deps/onnxruntime_gpu-1.14.0-cp38-cp38-linux_x86_64.whl" ]; then \
 	echo "Custom wheel exists, installing!" && \
-	pip install --force-reinstall /ai2modulus/deps/onnxruntime_gpu-1.14.0-cp38-cp38-linux_x86_64.whl; \
+	pip install --force-reinstall /modulus/deps/onnxruntime_gpu-1.14.0-cp38-cp38-linux_x86_64.whl; \
     else \
 	echo "No custom wheel present, skipping" && \
 	pip install numpy==1.22.4; \
     fi
 # cleanup of stage
-RUN rm -rf /ai2modulus/ 
+RUN rm -rf /modulus/ 
 
 # CI image
 FROM builder as ci
@@ -63,17 +63,17 @@ RUN mkdir -p /opt/cugraphops &&\
 
 ENV PYTHONPATH="${PYTHONPATH}:/opt/cugraphops/lib/python3.8/site-packages"
 
-COPY . /ai2modulus/
-RUN cd /ai2modulus/ && pip install -e . && rm -rf /ai2modulus/
+COPY . /modulus/
+RUN cd /modulus/ && pip install -e . && rm -rf /modulus/
 
 # Deployment image
 FROM builder as deploy
 RUN pip install protobuf==3.20.0 
-COPY . /ai2modulus/
-RUN cd /ai2modulus/ && pip install .
+COPY . /modulus/
+RUN cd /modulus/ && pip install .
 
 # Clean up
-RUN rm -rf /ai2modulus/ 
+RUN rm -rf /modulus/ 
 
 # Docs image
 FROM deploy as docs

Makefile

@@ -12,7 +12,7 @@ black:
 	black --check --exclude=docs/ ./
 
 interrogate:
-	cd ai2modulus && \
+	cd modulus && \
                 interrogate --ignore-init-method \
                 --ignore-init-module \
                 --ignore-module \
@@ -41,7 +41,7 @@ doctest:
 	coverage run \
                 --rcfile='test/coverage.docstring.rc' \
                 -m pytest \
-                --doctest-modules ai2modulus/ --ignore-glob=*internal*
+                --doctest-modules modulus/ --ignore-glob=*internal*
 
 pytest: 
 	coverage run \
@@ -59,11 +59,11 @@ coverage:
 		coverage html
 
 container-deploy:
-	docker build -t ai2modulus:deploy --target deploy -f Dockerfile .
+	docker build -t modulus:deploy --target deploy -f Dockerfile .
 
 container-ci:
-	docker build -t ai2modulus:ci --target ci -f Dockerfile .
+	docker build -t modulus:ci --target ci -f Dockerfile .
 
 container-docs:
-	docker build -t ai2modulus:docs --target docs -f Dockerfile .
+	docker build -t modulus:docs --target docs -f Dockerfile .
 

docs/conf.py

@@ -21,7 +21,7 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 import os
 import sphinx_rtd_theme
-from ai2modulus import __version__ as version
+from modulus import __version__ as version
 
 project = "NVIDIA Modulus"
 copyright = "2023, NVIDIA Modulus Team"

ai2modulus/datapipes/climate/sfno/dataloader.py → modulus/datapipes/climate/sfno/dataloader.py

@@ -20,7 +20,7 @@
 
 # distributed stuff
 import torch.distributed as dist
-from ai2modulus.utils.sfno.distributed import comm
+from modulus.utils.sfno.distributed import comm
 
 
 def init_distributed_io(params):  # pragma: no cover

ai2modulus/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py → modulus/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py

@@ -29,7 +29,7 @@
 
 # we need this for the zenith angle feature
 import datetime
-from ai2modulus.utils.sfno.zenith_angle import cos_zenith_angle
+from modulus.utils.sfno.zenith_angle import cos_zenith_angle
 
 
 class GeneralES(object):

ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py → modulus/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py

@@ -16,7 +16,7 @@
 import numpy as np
 
 # distributed stuff
-from ai2modulus.utils.sfno.distributed import comm
+from modulus.utils.sfno.distributed import comm
 
 # DALI stuff
 from nvidia.dali.pipeline import Pipeline
@@ -25,7 +25,7 @@
 from nvidia.dali.plugin.pytorch import DALIGenericIterator, LastBatchPolicy
 
 # es helper
-import ai2modulus.datapipes.climate.sfno.dataloaders.dali_es_helper_2d as esh
+import modulus.datapipes.climate.sfno.dataloaders.dali_es_helper_2d as esh
 
 
 class ERA5DaliESDataloader(object):

ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_dummy.py → modulus/datapipes/climate/sfno/dataloaders/data_loader_dummy.py

@@ -19,7 +19,7 @@
 import h5py
 
 # distributed stuff
-from ai2modulus.utils.sfno.distributed import comm
+from modulus.utils.sfno.distributed import comm
 
 
 class DummyLoader(object):

ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py → modulus/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py

@@ -20,7 +20,7 @@
 import h5py
 
 # import cv2
-from ai2modulus.utils.sfno.img_utils import reshape_fields
+from modulus.utils.sfno.img_utils import reshape_fields
 
 
 class MultifilesDataset(Dataset):

ai2modulus/datapipes/datapipe.py → modulus/datapipes/datapipe.py

@@ -17,7 +17,7 @@
 
 from typing import Union
 from pathlib import Path
-from ai2modulus.datapipes.meta import DatapipeMetaData
+from modulus.datapipes.meta import DatapipeMetaData
 
 
 class Datapipe:

ai2modulus/metrics/climate/acc.py → modulus/metrics/climate/acc.py

@@ -15,7 +15,7 @@
 # TODO(Dallas) Introduce Distributed Class for computation.
 
 import torch
-from ai2modulus.metrics.climate.reduction import _compute_lat_weights
+from modulus.metrics.climate.reduction import _compute_lat_weights
 
 Tensor = torch.Tensor
 

ai2modulus/metrics/climate/efi.py → modulus/metrics/climate/efi.py

@@ -15,8 +15,8 @@
 # TODO(Dallas) Introduce Distributed Class for computation.
 
 import torch
-from ai2modulus.metrics.general.histogram import normal_pdf, normal_cdf, histogram
-from ai2modulus.metrics.general.entropy import _entropy_from_counts
+from modulus.metrics.general.histogram import normal_pdf, normal_cdf, histogram
+from modulus.metrics.general.entropy import _entropy_from_counts
 
 Tensor = torch.Tensor
 

ai2modulus/metrics/climate/reduction.py → modulus/metrics/climate/reduction.py

@@ -15,7 +15,7 @@
 import torch
 from torch import Tensor
 from abc import ABC
-from ai2modulus.metrics.general.reduction import WeightedMean, WeightedVariance
+from modulus.metrics.general.reduction import WeightedMean, WeightedVariance
 
 
 def _compute_lat_weights(lat: Tensor) -> Tensor:

ai2modulus/metrics/general/calibration.py → modulus/metrics/general/calibration.py

@@ -14,7 +14,7 @@
 
 import torch
 import numpy as np
-from ai2modulus.metrics.general.histogram import histogram, linspace
+from modulus.metrics.general.histogram import histogram, linspace
 from typing import Union
 
 Tensor = torch.Tensor

ai2modulus/metrics/general/ensemble_metrics.py → modulus/metrics/general/ensemble_metrics.py

@@ -16,7 +16,7 @@
 import torch.distributed as dist
 from typing import Union, Tuple, List
 from abc import ABC, abstractmethod
-from ai2modulus.distributed.manager import DistributedManager
+from modulus.distributed.manager import DistributedManager
 from warnings import warn
 
 Tensor = torch.Tensor

ai2modulus/metrics/general/histogram.py → modulus/metrics/general/histogram.py

@@ -17,7 +17,7 @@
 import numpy
 import torch.distributed as dist
 from typing import Union, Tuple
-from ai2modulus.distributed.manager import DistributedManager
+from modulus.distributed.manager import DistributedManager
 from .ensemble_metrics import EnsembleMetrics
 
 Tensor = torch.Tensor

ai2modulus/models/afno/afno.py → modulus/models/afno/afno.py

@@ -15,8 +15,8 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import ai2modulus
-import ai2modulus.models.layers.fft as fft
+import modulus
+import modulus.models.layers.fft as fft
 
 from functools import partial
 from typing import Tuple, Any

ai2modulus/models/fno/fno.py → modulus/models/fno/fno.py

@@ -15,7 +15,8 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import ai2modulus.models.layers as layers
+import modulus.models.layers as layers
+import modulus
 
 from typing import Dict, List, Union, Tuple
 from torch import Tensor
@@ -484,7 +485,7 @@ class FNO(Module):
 
     Parameters
     ----------
-    decoder_net : ai2modulus.Module
+    decoder_net : modulus.Module
         Pointwise decoder network, input feature size should match `latent_channels`
     in_channels : int
         Number of input channels
@@ -508,14 +509,14 @@ class FNO(Module):
     Example
     -------
     >>> # define the decoder net
-    >>> decoder = ai2modulus.models.mlp.FullyConnected(
+    >>> decoder = modulus.models.mlp.FullyConnected(
     ...     in_features=32,
     ...     out_features=3,
     ...     num_layers=2,
     ...     layer_size=16,
     ... )
     >>> # define the 2d FNO model
-    >>> model = ai2modulus.models.fno.FNO(
+    >>> model = modulus.models.fno.FNO(
     ...     decoder_net=decoder,
     ...     in_channels=4,
     ...     dimension=2,

ai2modulus/models/gnn_layers/mesh_graph_mlp.py → modulus/models/gnn_layers/mesh_graph_mlp.py

@@ -22,7 +22,7 @@
 from torch.autograd.function import once_differentiable
 
 from .utils import concat_efeat, sum_efeat, CuGraphCSC
-from ai2modulus.models.layers.fused_silu import silu_backward_for
+from modulus.models.layers.fused_silu import silu_backward_for
 
 try:
     from apex.normalization import FusedLayerNorm

ai2modulus/models/graphcast/graph_cast_net.py → modulus/models/graphcast/graph_cast_net.py

@@ -19,18 +19,18 @@
 from typing import Any
 from dataclasses import dataclass
 
-from ai2modulus.models.gnn_layers.utils import set_checkpoint_fn, CuGraphCSC
-from ai2modulus.models.gnn_layers.embedder import (
+from modulus.models.gnn_layers.utils import set_checkpoint_fn, CuGraphCSC
+from modulus.models.gnn_layers.embedder import (
     GraphCastEncoderEmbedder,
     GraphCastDecoderEmbedder,
 )
-from ai2modulus.models.gnn_layers.mesh_graph_encoder import MeshGraphEncoder
-from ai2modulus.models.gnn_layers.mesh_graph_decoder import MeshGraphDecoder
-from ai2modulus.models.gnn_layers.mesh_graph_mlp import MeshGraphMLP
-from ai2modulus.models.module import Module
-from ai2modulus.models.meta import ModelMetaData
-from ai2modulus.utils.graphcast.graph import Graph
-from ai2modulus.utils.graphcast.data_utils import StaticData
+from modulus.models.gnn_layers.mesh_graph_encoder import MeshGraphEncoder
+from modulus.models.gnn_layers.mesh_graph_decoder import MeshGraphDecoder
+from modulus.models.gnn_layers.mesh_graph_mlp import MeshGraphMLP
+from modulus.models.module import Module
+from modulus.models.meta import ModelMetaData
+from modulus.utils.graphcast.graph import Graph
+from modulus.utils.graphcast.data_utils import StaticData
 
 from .graph_cast_processor import GraphCastProcessor
 

ai2modulus/models/graphcast/graph_cast_processor.py → modulus/models/graphcast/graph_cast_processor.py

@@ -19,9 +19,9 @@
 from torch import Tensor
 from dgl import DGLGraph
 
-from ai2modulus.models.gnn_layers.utils import set_checkpoint_fn, CuGraphCSC
-from ai2modulus.models.gnn_layers.mesh_node_block import MeshNodeBlock
-from ai2modulus.models.gnn_layers.mesh_edge_block import MeshEdgeBlock
+from modulus.models.gnn_layers.utils import set_checkpoint_fn, CuGraphCSC
+from modulus.models.gnn_layers.mesh_node_block import MeshNodeBlock
+from modulus.models.gnn_layers.mesh_edge_block import MeshEdgeBlock
 
 
 class GraphCastProcessor(nn.Module):

ai2modulus/models/layers/activations.py → modulus/models/layers/activations.py

@@ -14,7 +14,7 @@
 
 import torch
 import torch.nn as nn
-import ai2modulus
+import modulus
 
 Tensor = torch.Tensor
 

ai2modulus/models/layers/weight_norm.py → modulus/models/layers/weight_norm.py

@@ -15,7 +15,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import ai2modulus
+import modulus
 
 Tensor = torch.Tensor
 

ai2modulus/models/meshgraphnet/meshgraphnet.py → modulus/models/meshgraphnet/meshgraphnet.py

@@ -15,7 +15,7 @@
 import torch
 from torch import Tensor
 import torch.nn as nn
-import ai2modulus
+import modulus
 
 try:
     import dgl
@@ -29,13 +29,13 @@
 from typing import Union, List
 from dataclasses import dataclass
 
-from ai2modulus.models.meta import ModelMetaData
-from ai2modulus.models.module import Module
+from modulus.models.meta import ModelMetaData
+from modulus.models.module import Module
 
-from ai2modulus.models.gnn_layers.utils import CuGraphCSC
-from ai2modulus.models.gnn_layers.mesh_graph_mlp import MeshGraphMLP
-from ai2modulus.models.gnn_layers.mesh_edge_block import MeshEdgeBlock
-from ai2modulus.models.gnn_layers.mesh_node_block import MeshNodeBlock
+from modulus.models.gnn_layers.utils import CuGraphCSC
+from modulus.models.gnn_layers.mesh_graph_mlp import MeshGraphMLP
+from modulus.models.gnn_layers.mesh_edge_block import MeshEdgeBlock
+from modulus.models.gnn_layers.mesh_node_block import MeshNodeBlock
 
 
 @dataclass

ai2modulus/models/mlp/fully_connected.py → modulus/models/mlp/fully_connected.py

@@ -14,12 +14,12 @@
 
 import torch
 import torch.nn as nn
-import ai2modulus
+import modulus
 
 from torch import Tensor
 from dataclasses import dataclass
 from typing import Optional, Union, List
-from ai2modulus.models.layers import FCLayer
+from modulus.models.layers import FCLayer
 from ..meta import ModelMetaData
 from ..module import Module
 

ai2modulus/models/module.py → modulus/models/module.py

@@ -17,7 +17,7 @@
 
 from typing import Union
 from pathlib import Path
-from ai2modulus.models.meta import ModelMetaData
+from modulus.models.meta import ModelMetaData
 
 
 class Module(torch.nn.Module):

ai2modulus/models/pix2pix/pix2pix.py → modulus/models/pix2pix/pix2pix.py

@@ -58,7 +58,7 @@
 from typing import Union, List
 from dataclasses import dataclass
 
-import ai2modulus
+import modulus
 from ..meta import ModelMetaData
 from ..module import Module