Analysis model proposal

docs/userguide/components/analysis.md

@@ -0,0 +1,110 @@
+(analysis_model_userguide)=
+
+# Analysis Models
+
+Analysis models in Earth2Studio provides a set of models that serve as inference sinks
+and pipeline endpoints.
+Unlike {ref}`diagnostic_model_userguide` and {ref}`prognostic_model_userguide` which
+transform and return data, analysis models consume data without returning modified
+tensors or coordinates.
+
+Analysis models are designed for terminal operations in inference pipelines, such as:
+
+- Connections to downstream analysis tools / processes
+
+These models differ fundamentally from diagnostic and prognostic models as they do not
+produce data for further pipeline stages, instead serving as endpoints that process,
+analyze, or consume the input data.
+
+The list of analysis models that are already built into Earth2studio can be found in the
+API documentation {ref}`earth2studio.models.ax`.
+
+## Analysis Interface
+
+The full requirements for a standard analysis model are defined explicitly in the
+`earth2studio/models/ax/base.py`.
+
+```{literalinclude} ../../../earth2studio/models/ax/base.py
+:lines: 25-
+:language: python
+```
+
+:::{note}
+Analysis models do not need to inherit this protocol, this is simply used to define the
+required APIs.
+The key distinguishing feature is that the `__call__` method returns `None`, making
+these models inference sinks rather than data sources or transformers.
+:::
+
+## Analysis Usage
+
+### Execution
+
+Similar to diagnostic models, analysis models implement a {func}`__call__` function
+which takes in a data tensor with coordinate system and performs analysis operations
+without returning anything.
+
+```python
+# Assume model is an instance of an AnalysisModel
+x = torch.Tensor(...)  # Input tensor
+coords = CoordSystem(...)  # Coordinate system
+model(x, coords)  # Perform analysis - returns None
+```
+
+Since analysis models return `None`, they are typically used at the end of inference
+pipelines or in conjunction with other models:
+
+```python
+# Example: Analysis model used after prognostic prediction
+from earth2studio.models.px import PrognosticModel
+from earth2studio.models.ax import ValidationModel
+
+# Load models
+prognostic = PrognosticModel.load_model(prognostic_package)
+validator = ValidationModel.load_model(validation_package)
+
+# Generate prediction and analyze
+x, coords = prognostic(input_tensor, input_coords)
+validator(x, coords)  # Validate prediction (no return value)
+```
+
+## Custom Analysis Models
+
+Integrating your own analysis model is easy, just satisfy the interface above.
+The key requirement is that the `__call__` method returns `None` and performs whatever
+analysis, validation, or downstream processing operations are needed.
+
+We recommend users have a look at the {ref}`extension_examples` examples, which will
+step users through the simple process of implementing their own analysis model.
+
+```python
+class CustomAnalysisModel:
+    def __init__(self):
+        # Initialize your analysis model
+        pass
+
+    def __call__(self, x: torch.Tensor, coords: CoordSystem) -> None:
+        # Perform analysis operations
+        # Save results, compute metrics, validate data, etc.
+        # No return statement needed (returns None implicitly)
+        pass
+
+    def input_coords(self) -> CoordSystem:
+        # Return expected input coordinate system
+        return CoordSystem(...)
+
+    def to(self, device):
+        # Move model to device if needed
+        return self
+```
+
+## Contributing Analysis Models
+
+Want to add your analysis model to the package? Great, we will be happy to work with
+you.
+At the minimum we expect the model to abide by the defined interface and meet the
+requirements set forth in our contribution guide.
+
+Open an issue when you have an initial implementation you would like us to review.
+If you're aware of an existing analysis model and want us to implement it, open a
+feature request and we will get it triaged.

docs/userguide/index.md

@@ -92,6 +92,7 @@ about/data
 
 components/prognostic
 components/diagnostic
+components/analysis
 components/datasources
 components/perturbation
 components/io

earth2studio/models/ax/base.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+import torch
+
+from earth2studio.utils.type import CoordSystem
+
+
+@runtime_checkable
+class AnalysisModel(Protocol):
+    """Analysis model interface"""
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        coords: CoordSystem,
+    ) -> None:
+        """Execution of the analysis model that transforms physical data
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor intended to apply diagnostic function on
+        coords : CoordSystem
+            Ordered dict representing coordinate system that describes the tensor
+        """
+        pass
+
+    def input_coords(self) -> CoordSystem:
+        """Input coordinate system of analysis model
+
+        Returns
+        -------
+        CoordSystem
+            Coordinate system dictionary
+        """
+        pass
+
+    def to(self, device: Any) -> AnalysisModel:
+        """Moves analysis model onto inference device, this is typically satisfied via
+        `torch.nn.Module`. If applicable, otherwise this should be implemented as a
+        pass through.
+
+        Parameters
+        ----------
+        device : Any
+            Object representing the inference device, typically `torch.device` or str
+
+        Returns
+        -------
+        AnalysisModel
+            Returns instance of diagnostic
+        """
+        pass