Plugin System

Modern scientific and engineering workflows require flexible simulation frameworks that can be easily extended and customized. NeoFOAM’s plugin architecture is designed to enable users and developers to add new physics models, boundary conditions, and solver modules without modifying the core codebase. This approach promotes maintainability, collaboration, and rapid prototyping of new features.

Overview

The plugin system enables flexible model configuration via configuration files. By utilizing Pydantic, the model configuration is decoupled from the specific file format, allowing support for multiple formats such as OpenFOAM dictionaries, YAML, TOML, or JSON. It is a runtime-extensible system built on Pydantic discriminated unions and a registry pattern.

// RASModel acts as the discriminator
// This tag/type field in the config file
// selects the correct turbulence model
RASModel        kOmegaSST;

turbulence      on;

printCoeffs     on;

A discriminated union selects the appropriate configuration class based on a type/tag field (e.g., “RASModel”: “kOmegaSST” vs “RASModel”: “kEpsilon”). Each base class maintains its own registry of registered child implementations. All plugins and models are registered in a central registry, enabling easy access and management for UIs, validation, or generative AI integration. This allows users to retrieve all available plugins (such as turbulence models or boundary conditions), inspect their configuration options, and validate them. Additionally, the plugin system supports the generation of JSON schemas for documentation and validation.

Discriminated Unions in Pydantic

The discriminated union functionality in Pydantic is illustrated in the following simple example.

from typing import Literal, Union
from pydantic import BaseModel, Field, ValidationError

class Cat(BaseModel):
    # discriminator='pet_type'
    pet_type: Literal['cat']
    meows: int

class Dog(BaseModel):
    pet_type: Literal['dog']
    barks: float

class Lizard(BaseModel):
    pet_type: Literal['reptile', 'lizard']
    scales: bool

class Model(BaseModel):
    # pet is the discriminator_variable
    pet: Union[Cat, Dog, Lizard] = Field(..., discriminator='pet_type')
    n: int

print(Model(pet={'pet_type': 'dog', 'barks': 3.14}, n=1))
#> pet=Dog(pet_type='dog', barks=3.14) n=1

Discriminated unions allow instantiating the correct subclass based on the value of a discriminator field (here pet_type). The discriminator_variable (here pet) is the field that accepts the union of possible types. Therefore, all available models must be registered in the Union to be selectable via the discriminator. This facilitates easy validation of configuration files, as all required information is stored directly in the Pydantic models.

Plugin System and Registration of Subclasses

The PluginSystem stores and registers all base classes (similar to the Model shown above) in a centralized registry. This ensures all plugins are easily accessible and usable by other parts of the codebase or external tools.

The main challenge lies in dynamically creating the Union type. This is required because the plugin system enables the registration of new subclasses at runtime.

The example below demonstrates how to register a new base class with two subclasses using the PluginSystem. To automatically create the discriminated union, the base class must be decorated with @PluginSystem.register(), providing the discriminator_variable and discriminator. This automatically adds the discriminated union field to the base class, similar to the Pydantic example above.

from pydantic import BaseModel
from typing import Literal
from neofoam.core.plugin_system import PluginSystem

@PluginSystem.register(discriminator_variable="shape", discriminator="shape_type")
class ShapeInterface(BaseModel):
    # the decorator automatically adds shape as discriminator variable
    # and shape_type as discriminator
    # the commented line below is automatically added:
    # shape: Union[CircleConfig, SquareConfig] = Field(discriminator='shape_type')
    color: str
    name: str = "default"

@ShapeInterface.register
class CircleConfig(BaseModel):
    shape_type: Literal["circle"]
    radius: float

@ShapeInterface.register
class SquareConfig(BaseModel):
    shape_type: Literal["square"]
    side: float

The ShapeInterface class now contains a dynamically created field shape, which is a discriminated union of all registered subclasses. The subclasses CircleConfig and SquareConfig are registered using the @ShapeInterface.register decorator. This action automatically updates the union type and recreates the model to include the new subclasses.

The model must be instantiated using the class method create() to ensure that the latest version of the union type is used, as Pydantic models cache their schema definitions.

# return a circle instance
circle = ShapeInterface.create(
    shape={"shape_type": "circle", "radius": 5.0}, # shows the ability to serialize via discriminator
    color="red"
)

# return a square instance
square = ShapeInterface.create(
    shape={"shape_type": "square", "side": 3.0},
    color="blue"
)

# Accessing plugin-specific fields
assert circle.shape.radius == 5.0
assert square.shape.side == 3.0

JSON Schema Generation

The dynamically generated plugin models support JSON schema generation for validation and documentation. JSON schemas expose all possible configurations for the available plugins of each base class, enabling the creation of user interfaces or configuration file validators.

ShapeInterface.plugin_model provides access to the dynamically updated model and is required to generate up-to-date schemas.

# Get JSON schema for validation and documentation
Shape = ShapeInterface.plugin_model
schema = Shape.model_json_schema()

# Schema includes discriminator mapping
discriminator = schema["properties"]["shape"]["discriminator"]
print(discriminator["propertyName"])  # "shape_type"
print(discriminator["mapping"])       # {"circle": "...", "square": "..."}

Registry Management

The PluginSystem class provides utility methods to manage and inspect the plugin registries. The following example demonstrates how to list registered plugins, retrieve specific registries, and remove plugins.

# List all plugins
all_plugins = PluginSystem.list_plugins()
print(all_plugins.keys())  # ["turbulence models", "thermodynamic models", ...]

# Get specific registry
shape_registry = PluginSystem.get_registered("TurbulenceModel")
plugin_classes = shape_registry.plugin_registry

# Remove plugins
success = PluginSystem.remove_plugin_model("TurbulenceModel", KOmegaSSTModel)

Registration via Entrypoint

Warning

This feature is not yet implemented.