Skip to content

Extending Models

This document provides detailed information on how to add new ocean models to Rompy. For basic concepts about model configuration, please see the User Guide and Progressive Tutorials. For advanced architecture information, see Plugin Architecture.

Overview

Rompy is designed to be extensible, allowing new ocean models to be integrated through its plugin architecture. Adding a new model typically involves:

  1. Creating a model-specific configuration class
  2. Defining the model's template files
  3. Implementing any model-specific functionality
  4. Registering the model with Rompy's plugin system

Creating a New Model Configuration

1. Basic Model Configuration

All model configurations in Rompy inherit from Pydantic models. Here's how to create a basic configuration for a new model:

from typing import Literal
from pydantic import Field
from rompy.core.config import BaseConfig
from rompy.core.grid import BaseGrid

class MyNewModelConfig(BaseConfig):
    """Configuration for MyNewModel."""

    # Required: Model type identifier for discriminated unions
    model_type: Literal["mynewmodel"] = Field(
        "mynewmodel", 
        description="Type of the model, used for discriminated union"
    )

    # Model-specific parameters
    grid: BaseGrid = Field(
        ..., 
        description="Grid configuration for MyNewModel"
    )

    # Add your model-specific parameters here
    physics_options: dict = Field(
        default_factory=dict,
        description="Physics options for MyNewModel"
    )

    # Template definition (if using templating)
    template: str = Field(
        "mynewmodel/templates/", 
        description="Path to model templates"
    )

    class Config:
        extra = "forbid"  # Forbid extra fields not defined in the model

2. Model-Specific Data Handling

If your model requires specific data handling capabilities, you can create custom data classes:

from rompy.core.data import DataGrid
from rompy.core.source import SourceBase

class MyNewModelData(DataGrid):
    """Specialized data class for MyNewModel."""

    model_type: Literal["mynewmodeldata"] = Field(
        "mynewmodeldata", 
        description="Type of the data source"
    )

    # Model-specific data properties
    coordinate_system: str = Field(
        "cartesian",
        description="Coordinate system for the data"
    )

    # Custom validation methods if needed
    def get_coordinates(self):
        """Custom method to extract coordinates in model-specific format."""
        # Implementation specific to your model
        pass

3. Model-Specific Grid Configuration

If your model has special grid requirements, create a specialized grid class:

from rompy.core.grid import BaseGrid
from typing import Literal
from pydantic import Field

class MyNewModelGrid(BaseGrid):
    """Specialized grid for MyNewModel."""

    model_type: Literal["mynewmodelgrid"] = Field(
        "mynewmodelgrid", 
        description="Type of the grid"
    )

    # Grid-specific parameters for your model
    grid_type: str = Field(
        "structured", 
        description="Type of grid (structured, unstructured, etc.)"
    )

    def to_model_format(self):
        """Convert the grid to the format required by MyNewModel."""
        # Implementation to convert to your model's grid format
        pass

Model Templates

Most ocean models use configuration files with specific formats. Rompy uses templating to generate these files from your configuration.

1. Template Directory Structure

Create a directory structure for your model's templates:

mynewmodel/
├── templates/
│   ├── model_config.j2     # Main configuration file
│   ├── grid_def.j2         # Grid definition
│   ├── forcing.j2          # Forcing data specification
│   └── output.j2           # Output configuration

2. Template Example

Here's an example template for a model configuration file:

# MyNewModel Configuration File
# Generated by Rompy

# Time settings
start_time = {{ runtime.period.start.strftime('%Y-%m-%d %H:%M:%S') }}
end_time = {{ runtime.period.end.strftime('%Y-%m-%d %H:%M:%S') }}
time_step = {{ config.time_step_seconds }}s

# Grid settings
grid_file = {{ config.grid.to_model_format() }}

# Physics options
{%- for option, value in config.physics_options.items() %}
{{ option }} = {{ value }}
{%- endfor %}

# Output settings
output_dir = {{ runtime.output_dir }}/{{ runtime.run_id }}
output_interval = {{ config.output_interval_seconds }}s

3. Template Integration

Integrate your templates with your model configuration:

class MyNewModelConfig(BaseConfig):
    model_type: Literal["mynewmodel"] = "mynewmodel"

    # Configuration parameters...
    time_step_seconds: int = Field(
        300,  # 5 minutes default
        description="Time step in seconds"
    )

    output_interval_seconds: int = Field(
        3600,  # 1 hour default
        description="Output interval in seconds"
    )

    # Template settings for this model
    template: str = Field(
        "mynewmodel/templates/model_config.j2",
        description="Path to the main configuration template"
    )

    def get_model_variables(self, run) -> dict:
        """Get variables specific to this model for template rendering."""
        variables = super().get_render_variables(run)
        variables.update({
            "time_step_seconds": self.time_step_seconds,
            "output_interval_seconds": self.output_interval_seconds,
        })
        return variables

Registering Your Model

To make your new model available in Rompy, you need to register it using entry points.

1. Update pyproject.toml

Add your model configuration to the rompy.config entry point:

[project.entry-points."rompy.config"]
mynewmodel = "mynewmodel.config:MyNewModelConfig"

2. Entry Points Overview

Rompy uses several entry point groups:

  • rompy.config: Model configuration classes
  • rompy.run: Execution backend implementations
  • rompy.postprocess: Post-processing implementations
  • rompy.source: Data source implementations

Complete Model Extension Example

Here's a complete example of a new model extension:

# mynewmodel/config.py
from typing import Literal, Optional
from pydantic import Field, BaseModel
from rompy.core.config import BaseConfig
from rompy.core.grid import BaseGrid

class MyNewModelPhysics(BaseModel):
    """Physics configuration for MyNewModel."""
    baroclinic: bool = Field(
        True,
        description="Whether to include baroclinic effects"
    )
    viscosity: float = Field(
        0.01,
        description="Horizontal viscosity coefficient"
    )
    diffusion: float = Field(
        0.001,
        description="Horizontal diffusion coefficient"
    )

class MyNewModelConfig(BaseConfig):
    """Complete configuration for MyNewModel."""

    model_type: Literal["mynewmodel"] = Field("mynewmodel", description="Model identifier")

    # Core configuration
    grid: BaseGrid = Field(..., description="Model grid")
    physics: MyNewModelPhysics = Field(
        default_factory=MyNewModelPhysics,
        description="Physics configuration"
    )

    # Model-specific parameters
    time_step_seconds: int = Field(
        300,
        description="Model time step in seconds"
    )

    # Template reference
    template: str = Field(
        "mynewmodel/templates/main.j2",
        description="Main configuration template"
    )

    def get_model_variables(self, run) -> dict:
        """Get variables specific to this model for template rendering."""
        return {
            "physics": self.physics.dict(),
            "time_step": self.time_step_seconds,
            "grid_file": self.grid.to_model_format(),
        }

# In pyproject.toml
[project.entry-points."rompy.config"]
mynewmodel = "mynewmodel.config:MyNewModelConfig"

Model Testing

When adding a new model, ensure you include comprehensive tests:

1. Unit Tests

Test the configuration validation:

import pytest
from mynewmodel.config import MyNewModelConfig
from rompy.core.grid import RegularGrid

def test_mynewmodel_config_creation():
    """Test creating a valid MyNewModel configuration."""
    grid = RegularGrid(
        lon_min=-75.0, lon_max=-65.0,
        lat_min=35.0, lat_max=45.0,
        dx=0.1, dy=0.1
    )

    config = MyNewModelConfig(
        grid=grid,
        physics=MyNewModelPhysics(baroclinic=False)
    )

    assert config.model_type == "mynewmodel"
    assert config.physics.baroclinic == False

2. Integration Tests

Test that your model works with the full Rompy workflow:

from rompy.model import ModelRun
from rompy.core.time import TimeRange
from datetime import datetime

def test_mynewmodel_full_workflow():
    """Test full workflow with MyNewModel."""
    # Assuming you have a valid MyNewModelConfig instance
    config = MyNewModelConfig(
        # ... configuration parameters
    )

    run = ModelRun(
        run_id="test_mynewmodel",
        period=TimeRange(
            start=datetime(2023, 1, 1),
            end=datetime(2023, 1, 2),
            interval="1H"
        ),
        config=config,
        output_dir="./test_output"
    )

    # Test generation phase
    staging_dir = run.generate()
    assert staging_dir.exists()

    # Note: Execution tests may require the actual model to be installed
    # This is typically done in CI/CD with appropriate test fixtures

Best Practices

1. Validation

Implement custom validation for model-specific constraints:

from pydantic import validator

class MyNewModelConfig(BaseConfig):
    # ... other fields

    time_step_seconds: int = Field(300, description="Time step in seconds")
    depth_min: float = Field(0.1, description="Minimum depth in meters")

    @validator('time_step_seconds')
    def validate_time_step(cls, v):
        if v <= 0:
            raise ValueError('Time step must be positive')
        if v > 3600:  # More than 1 hour
            raise ValueError('Time step seems too large')
        return v

    @validator('depth_min')
    def validate_depth_min(cls, v):
        if v <= 0:
            raise ValueError('Minimum depth must be positive')
        return v

2. Default Values

Provide sensible default values for all parameters:

class MyNewModelConfig(BaseConfig):
    # ... other fields

    # Good: Sensible default
    output_format: str = Field(
        "netcdf",
        description="Output format"
    )

    # Good: Factory for complex defaults
    physics: MyNewModelPhysics = Field(
        default_factory=MyNewModelPhysics,
        description="Physics configuration"
    )

3. Documentation

Document all parameters and their expected ranges:

class MyNewModelConfig(BaseConfig):
    stability_factor: float = Field(
        0.5,
        description=(
            "Stability factor for numerical scheme. "
            "Typically between 0.1 and 0.9. Lower values are more stable but slower."
        ),
        ge=0.01,  # Greater than or equal to 0.01
        le=0.99,  # Less than or equal to 0.99
    )

Troubleshooting Common Issues

1. Plugin Registration

If your model doesn't appear in Rompy after registration:

  1. Verify entry point is correctly defined in pyproject.toml
  2. Reinstall the package: pip install -e .
  3. Check that the class path is correct (module:object)

2. Template Issues

If templates are not rendering correctly:

  1. Verify template paths in configuration
  2. Check that template syntax is valid Jinja2
  3. Ensure required variables are provided to the template

3. Validation Problems

If validation is failing unexpectedly:

  1. Check Pydantic version compatibility
  2. Validate the data being passed to the configuration
  3. Consider using .dict(exclude_unset=True) to see what values are set

Advanced Extension Features

1. Model-Specific Backends

If your model needs special execution handling, you can also create custom backends:

# mynewmodel/backends.py
from rompy.backends.config import BaseBackendConfig
from rompy.run import BaseRunBackend

class MyNewModelBackendConfig(BaseBackendConfig):
    """Backend config specific to MyNewModel."""
    model_type: Literal["mynewmodelbackend"] = "mynewmodelbackend"
    special_parameter: str = Field(
        "default_value",
        description="Model-specific backend parameter"
    )

class MyNewModelRunBackend(BaseRunBackend):
    """Custom execution logic for MyNewModel."""
    def run(self, model_run, config: MyNewModelBackendConfig):
        # Custom execution logic here
        # This might include model-specific preprocessing or postprocessing
        pass

2. Model-Specific Postprocessors

Create custom postprocessors for model-specific results:

# mynewmodel/postprocess.py
from rompy.postprocess import BasePostprocessor

class MyNewModelProcessor(BasePostprocessor):
    """Postprocessor specific to MyNewModel results."""
    def process(self, model_run, **kwargs):
        # Custom result processing for MyNewModel
        output_dir = f"{model_run.output_dir}/{model_run.run_id}"
        # Implement your processing logic here
        pass

Next Steps

  • Review the Plugin Architecture for more details on the extension system
  • Check the Architecture Overview for understanding how components interact
  • Follow the Developer Guide for development best practices
  • Look at existing model implementations in rompy-swan and rompy-schism for more examples
  • Test your model extension thoroughly before contributing