Adding a few introductory tutorials (#233)

* add a few beginner level tutorials

* make changes to exclude the docs folder from ci

Author: ktangsali

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
           "--ignore-magic", "--fail-under=99", "--exclude=['setup.py', 'test', 'build', 'docs']",  
           "--ignore-regex=['forward', 'backward', 'reset_parameters', 'extra_repr', 'MetaData', 'apply_activation','exec_activation']", 
           "--color", "--"]
-      exclude: ^modulus/internal/|^modulus/experimental/
+      exclude: ^modulus/internal/|^modulus/experimental/|^docs/
 
 -   repo: https://github.com/igorshubovych/markdownlint-cli
     rev: v0.35.0
@@ -49,6 +49,7 @@ repos:
     hooks:
     - id: ruff
       args: [--fix]
+      exclude: ^docs/
 
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v3.4.0

Makefile

@@ -35,7 +35,7 @@ doctest:
 pytest: 
 	coverage run \
 		--rcfile='test/coverage.pytest.rc' \
-		-m pytest 
+		-m pytest --ignore-glob=*docs* 
 
 pytest-internal:
 	cd test/internal && \

docs/api/modulus.models.rst

@@ -248,6 +248,8 @@ complex models.
     refactored in the future.
 
 
+.. _modulus-models-from-torch:
+
 Converting PyTorch Models to Modulus Models
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -284,6 +286,10 @@ below.
     ModulusModel = modulus.Module.from_torch(TorchModel, meta=ConvMetaData())
 
 
+
+
+.. _saving-and-loading-modulus-models:
+
 Saving and Loading Modulus Models
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -500,37 +506,3 @@ DLWP Model
     :members:
     :show-inheritance:
 
-Spherical Harmonics Fourier Neural Operator
---------------------------------------------
-
-.. automodule:: modulus.models.sfno.sfnonet
-    :members:
-    :show-inheritance:
-
-.. automodule:: modulus.models.sfno.activations
-    :members:
-    :show-inheritance:
-
-.. automodule:: modulus.models.sfno.factorizations
-    :members:
-    :show-inheritance:
-
-.. automodule:: modulus.models.sfno.layers
-    :members:
-    :show-inheritance:
-
-.. automodule:: modulus.models.sfno.s2convolutions
-    :members:
-    :show-inheritance:
-
-.. automodule:: modulus.models.sfno.contractions
-    :members:
-    :show-inheritance:
-
-.. automodule:: modulus.models.sfno.initialization
-    :members:
-    :show-inheritance:
-
-.. automodule:: modulus.models.sfno.preprocessor
-    :members:
-    :show-inheritance:

docs/index.rst

@@ -1,6 +1,13 @@
 Welcome to Modulus Core's documentation!
 ========================================
 
+.. toctree::
+   :maxdepth: 2
+   :caption: Modulus Tutorials
+   :name: Modulus Tutorials
+
+   tutorials/simple_training_example.rst
+   tutorials/simple_logging_and_checkpointing.rst
 
 .. toctree::
    :maxdepth: 2

docs/test_runner.py

@@ -0,0 +1,18 @@
+import os
+import subprocess
+
+import pytest
+
+# Collecting all the Python files in the scripts directory
+script_files = [f for f in os.listdir("test_scripts/") if f.endswith('.py')]
+
+#print(script_files)
+@pytest.mark.parametrize("script_file", script_files)
+def test_script_execution(script_file):
+    """Test if a script runs without error."""
+    filepath = os.path.join("test_scripts/", script_file)
+    print(filepath)
+    result = subprocess.run(["python", filepath], capture_output=True, text=True)
+
+    # Check that the script executed successfully
+    assert result.returncode == 0, f"Script {script_file} failed with error:\n{result.stderr}"

docs/test_scripts/test_basic.py

@@ -0,0 +1,47 @@
+# [imports]
+import torch
+
+import modulus
+from modulus.datapipes.benchmarks.darcy import Darcy2D
+from modulus.metrics.general.mse import mse
+from modulus.models.fno.fno import FNO
+
+# [imports]
+
+# [code]
+normaliser = {
+    "permeability": (1.25, 0.75),
+    "darcy": (4.52e-2, 2.79e-2),
+}
+dataloader = Darcy2D(
+    resolution=256, batch_size=64, nr_permeability_freq=5, normaliser=normaliser
+)
+model = FNO(
+    in_channels=1,
+    out_channels=1,
+    decoder_layers=1,
+    decoder_layer_size=32,
+    dimension=2,
+    latent_channels=32,
+    num_fno_layers=4,
+    num_fno_modes=12,
+    padding=5,
+).to("cuda")
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+scheduler = torch.optim.lr_scheduler.LambdaLR(
+    optimizer, lr_lambda=lambda step: 0.85**step
+)
+
+# run for 20 iterations
+for i in range(20):
+    batch = next(iter(dataloader))
+    true = batch["darcy"]
+    pred = model(batch["permeability"])
+    loss = mse(pred, true)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+
+    print(f"Iteration: {i}. Loss: {loss.detach().cpu().numpy()}")
+# [code]

docs/test_scripts/test_basic_checkpointing.py

@@ -0,0 +1,69 @@
+# [imports]
+import torch
+
+import modulus
+from modulus.datapipes.benchmarks.darcy import Darcy2D
+from modulus.launch.utils import load_checkpoint, save_checkpoint
+from modulus.metrics.general.mse import mse
+from modulus.models.fno.fno import FNO
+
+# [imports]
+
+# [code]
+normaliser = {
+    "permeability": (1.25, 0.75),
+    "darcy": (4.52e-2, 2.79e-2),
+}
+dataloader = Darcy2D(
+    resolution=256, batch_size=64, nr_permeability_freq=5, normaliser=normaliser
+)
+model = FNO(
+    in_channels=1,
+    out_channels=1,
+    decoder_layers=1,
+    decoder_layer_size=32,
+    dimension=2,
+    latent_channels=32,
+    num_fno_layers=4,
+    num_fno_modes=12,
+    padding=5,
+).to("cuda")
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+scheduler = torch.optim.lr_scheduler.LambdaLR(
+    optimizer, lr_lambda=lambda step: 0.85**step
+)
+
+# load the epoch and optimizer, model ans scheduler parameters from the checkpoint if
+# it exists
+loaded_epoch = load_checkpoint(
+    "./checkpoints",
+    models=model,
+    optimizer=optimizer,
+    scheduler=scheduler,
+    device="cuda",
+)
+
+# we will setup the training to run for 20 epochs each epoch running for 5 iterations
+# starting with the loaded epoch
+for i in range(max(1, loaded_epoch), 20):
+    # this would be iterations through different batches
+    for _ in range(5):
+        batch = next(iter(dataloader))
+        true = batch["darcy"]
+        pred = model(batch["permeability"])
+        loss = mse(pred, true)
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+
+    # save checkpoint every 5th epoch
+    if i % 5 == 0:
+        save_checkpoint(
+            "./checkpoints",
+            models=model,
+            optimizer=optimizer,
+            scheduler=scheduler,
+            epoch=i,
+        )
+# [code]

docs/test_scripts/test_basic_inference.py

@@ -0,0 +1,41 @@
+# [imports]
+import torch
+
+import modulus
+from modulus.models.fno.fno import FNO
+
+# [imports]
+
+# [code]
+
+model = FNO(
+    in_channels=1,
+    out_channels=1,
+    decoder_layers=1,
+    decoder_layer_size=32,
+    dimension=2,
+    latent_channels=32,
+    num_fno_layers=4,
+    num_fno_modes=12,
+    padding=5,
+).to("cuda")
+
+# Save the checkpoint. For demo, we will just save untrained checkpoint,
+# but in typical workflows is saved after model training.
+model.save("untrained_checkpoint.mdlus")
+
+# Inference code
+
+# The parameters to instantitate the model will be loaded from the checkpoint
+model_inf = modulus.Module.from_checkpoint("untrained_checkpoint.mdlus").to("cuda")
+
+# put the model in evaluation mode
+model_inf.eval()
+
+# run inference
+with torch.inference_mode():
+    input = torch.ones(8, 1, 256, 256).to("cuda")
+    output = model_inf(input)
+    print(output.shape)
+
+# [code]

docs/test_scripts/test_console_logger.py

@@ -0,0 +1,63 @@
+# [imports]
+import torch
+
+import modulus
+from modulus.datapipes.benchmarks.darcy import Darcy2D
+from modulus.launch.logging import LaunchLogger, PythonLogger
+from modulus.metrics.general.mse import mse
+from modulus.models.fno.fno import FNO
+
+# [imports]
+
+# [code]
+normaliser = {
+    "permeability": (1.25, 0.75),
+    "darcy": (4.52e-2, 2.79e-2),
+}
+dataloader = Darcy2D(
+    resolution=256, batch_size=64, nr_permeability_freq=5, normaliser=normaliser
+)
+model = FNO(
+    in_channels=1,
+    out_channels=1,
+    decoder_layers=1,
+    decoder_layer_size=32,
+    dimension=2,
+    latent_channels=32,
+    num_fno_layers=4,
+    num_fno_modes=12,
+    padding=5,
+).to("cuda")
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+scheduler = torch.optim.lr_scheduler.LambdaLR(
+    optimizer, lr_lambda=lambda step: 0.85**step
+)
+
+# Initialize the logger
+logger = PythonLogger("main")  # General python logger
+LaunchLogger.initialize()
+
+# Use logger methods to track various information during training
+logger.info("Starting Training!")
+
+# we will setup the training to run for 20 epochs each epoch running for 5 iterations
+for i in range(20):
+    # wrap the epoch in launch logger to control frequency of output for console logs
+    with LaunchLogger("train", epoch=i) as launchlog:
+        # this would be iterations through different batches
+        for _ in range(5):
+            batch = next(iter(dataloader))
+            true = batch["darcy"]
+            pred = model(batch["permeability"])
+            loss = mse(pred, true)
+            loss.backward()
+            optimizer.step()
+            scheduler.step()
+            launchlog.log_minibatch({"Loss": loss.detach().cpu().numpy()})
+
+        launchlog.log_epoch({"Learning Rate": optimizer.param_groups[0]["lr"]})
+
+
+logger.info("Finished Training!")
+# [code]