Platform API Reference¶

The opifex.platform package provides community platform infrastructure for storing, versioning, and sharing neural functionals and scientific ML models.

Overview¶

The platform module offers:

Neural Functional Registry: Store and version neural network weights and architectures
Model Search: Semantic search over registered models
Version Control: Track model evolution and lineage
Validation: Ensure model compatibility and correctness
Collaboration (planned): Team collaboration features
Dashboard (planned): Analytics and monitoring

Registry Core¶

RegistryService¶

Central registry for neural functionals with versioning and metadata.

from opifex.platform.registry import RegistryService

class RegistryService:
    """
    Registry for neural functionals with version control and metadata.

    Neural functionals are parameterized neural networks that solve
    specific scientific computing tasks (PDEs, operators, etc.).

    Args:
        storage_path: Path to registry storage directory
        enable_caching: Whether to cache frequently accessed models
        cache_size: Maximum number of cached models

    Attributes:
        models: Dictionary of registered models
        metadata_store: Model metadata database
        version_graph: DAG of model versions
    """

    def __init__(
        self,
        storage_path: str = "~/.opifex/registry",
        enable_caching: bool = True,
        cache_size: int = 100
    ):
        """Initialize registry with storage configuration."""

Methods¶

`register(model, metadata) -> str`¶

Register a new neural functional.

def register(
    self,
    model: nnx.Module,
    metadata: FunctionalMetadata,
    version: Optional[str] = None,
    parent_id: Optional[str] = None
) -> str:
    """
    Register neural functional in the registry.

    Args:
        model: Flax NNX model to register
        metadata: Model metadata (name, description, etc.)
        version: Version string (auto-generated if None)
        parent_id: ID of parent model (for versioning)

    Returns:
        Unique model ID

    Example:
        >>> from opifex.neural.operators.fno import FNO
        >>> model = FNO(modes=12, width=32)
        >>> metadata = FunctionalMetadata(
        ...     name="darcy-flow-fno",
        ...     description="FNO for Darcy flow prediction",
        ...     task="operator-learning",
        ...     domain="fluid-dynamics",
        ...     tags=["pde", "elliptic", "darcy"]
        ... )
        >>> model_id = registry.register(model, metadata)
        >>> print(f"Registered as: {model_id}")
    """

`load(model_id, version='latest') -> nnx.Module`¶

Load registered model from registry.

def load(
    self,
    model_id: str,
    version: str = "latest",
    device: Optional[str] = None
) -> nnx.Module:
    """
    Load model from registry.

    Args:
        model_id: Unique model identifier
        version: Version to load ('latest', specific version string)
        device: Target device ('cpu', 'cuda', 'tpu')

    Returns:
        Loaded Flax NNX model

    Example:
        >>> model = registry.load("darcy-flow-fno", version="latest")
        >>> # Use for inference
        >>> prediction = model(test_input)
    """

`update(model_id, model, metadata) -> str`¶

Update existing model with new version.

def update(
    self,
    model_id: str,
    model: nnx.Module,
    metadata: Optional[FunctionalMetadata] = None,
    version_note: str = ""
) -> str:
    """
    Create new version of existing model.

    Args:
        model_id: ID of model to update
        model: Updated model
        metadata: Updated metadata (None = keep existing)
        version_note: Description of changes

    Returns:
        New version ID

    Example:
        >>> # Fine-tune existing model
        >>> model = registry.load("darcy-flow-fno")
        >>> # ... training ...
        >>> new_version = registry.update(
        ...     "darcy-flow-fno",
        ...     model,
        ...     version_note="Fine-tuned on high-Reynolds data"
        ... )
    """

`search(query, filters) -> List[str]`¶

Search for models using semantic query and filters.

def search(
    self,
    query: str = "",
    filters: Optional[Dict[str, Any]] = None,
    limit: int = 10,
    sort_by: str = "relevance"
) -> List[SearchResult]:
    """
    Search registry for relevant models.

    Args:
        query: Free-text search query
        filters: Filter criteria:
            - 'task': Task type (operator-learning, pinn, dft)
            - 'domain': Physics domain (fluid-dynamics, quantum)
            - 'tags': List of required tags
            - 'min_accuracy': Minimum validation accuracy
        limit: Maximum number of results
        sort_by: Sort criterion (relevance, accuracy, date)

    Returns:
        List of search results with model IDs and metadata

    Example:
        >>> # Find fluid dynamics models
        >>> results = registry.search(
        ...     query="fluid flow prediction",
        ...     filters={
        ...         'domain': 'fluid-dynamics',
        ...         'task': 'operator-learning',
        ...         'min_accuracy': 0.95
        ...     },
        ...     limit=5
        ... )
        >>> for result in results:
        ...     print(f"{result.name}: {result.score:.2f}")
    """

`list_versions(model_id) -> List[str]`¶

List all versions of a model.

def list_versions(
    self,
    model_id: str,
    include_metadata: bool = True
) -> List[VersionInfo]:
    """
    Get version history for model.

    Args:
        model_id: Model identifier
        include_metadata: Include full metadata for each version

    Returns:
        List of version information

    Example:
        >>> versions = registry.list_versions("darcy-flow-fno")
        >>> for v in versions:
        ...     print(f"{v.version}: {v.timestamp} - {v.note}")
    """

`delete(model_id, version=None)`¶

Delete model or specific version.

def delete(
    self,
    model_id: str,
    version: Optional[str] = None,
    cascade: bool = False
) -> None:
    """
    Delete model or version from registry.

    Args:
        model_id: Model to delete
        version: Specific version (None = delete all versions)
        cascade: If True, also delete dependent models

    Example:
        >>> # Delete specific version
        >>> registry.delete("darcy-flow-fno", version="v0.1.0")
        >>>
        >>> # Delete entire model
        >>> registry.delete("old-model", cascade=True)
    """

Model Metadata¶

FunctionalMetadata¶

Structured metadata for neural functionals.

from opifex.platform.registry import FunctionalMetadata

@dataclass
class FunctionalMetadata:
    """
    Metadata for registered neural functionals.

    Attributes:
        name: Human-readable model name
        description: Detailed description
        task: Task type (operator-learning, pinn, neural-dft, etc.)
        domain: Physics domain (fluid-dynamics, quantum, etc.)
        tags: List of descriptive tags
        architecture: Architecture name (FNO, DeepONet, etc.)
        input_shape: Expected input shape
        output_shape: Expected output shape
        performance_metrics: Validation metrics
        training_config: Training configuration used
        created_at: Creation timestamp
        updated_at: Last update timestamp
        author: Model author/organization
        license: Software license
        paper_url: Link to associated paper
        code_url: Link to training code
    """

    name: str
    description: str
    task: str
    domain: str
    tags: List[str] = field(default_factory=list)
    architecture: Optional[str] = None
    input_shape: Optional[Tuple[int, ...]] = None
    output_shape: Optional[Tuple[int, ...]] = None
    performance_metrics: Dict[str, float] = field(default_factory=dict)
    training_config: Dict[str, Any] = field(default_factory=dict)
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None
    author: Optional[str] = None
    license: str = "MIT"
    paper_url: Optional[str] = None
    code_url: Optional[str] = None

Example Usage¶

metadata = FunctionalMetadata(
    name="burgers-fno-large",
    description="Large FNO trained on Burgers equation dataset",
    task="operator-learning",
    domain="fluid-dynamics",
    architecture="FNO",
    tags=["pde", "nonlinear", "burgers", "turbulence"],
    input_shape=(256,),
    output_shape=(100, 256),
    performance_metrics={
        'val_mse': 1.2e-4,
        'val_relative_l2': 0.023,
        'inference_time_ms': 5.2
    },
    training_config={
        'epochs': 500,
        'batch_size': 32,
        'learning_rate': 1e-3,
        'optimizer': 'adam'
    },
    author="Opifex Team",
    license="MIT",
    paper_url="https://arxiv.org/abs/2010.08895"
)

Model Search¶

SearchEngine¶

Search capabilities for discovering neural functionals in the registry.

from opifex.platform.registry import SearchEngine

opifex.platform.registry.search.SearchEngine ¶

SearchEngine(registry_service, enable_semantic_search: bool = True, similarity_threshold: float = 0.7)

Neural functional search engine.

Provides full search capabilities including text-based search, semantic search with neural embeddings, advanced filtering, and recommendation systems.

Parameters:

Name	Type	Description	Default
`registry_service`		Registry service for functional access	required
`enable_semantic_search`	`bool`	Whether to enable neural embeddings	`True`
`similarity_threshold`	`float`	Minimum similarity for semantic matches	`0.7`

search `async` ¶

search(query: SearchQuery) -> list[SearchResult]

Execute search query for neural functionals.

Parameters:

Name	Type	Description	Default
`query`	`SearchQuery`	Search query with criteria and options	required

Returns:

Type	Description
`list[SearchResult]`	List of matching functionals ranked by relevance

suggest_functionals `async` ¶

suggest_functionals(functional_id: str, limit: int = 5) -> list[SearchResult]

Suggest similar functionals based on a reference functional.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of reference functional	required
`limit`	`int`	Maximum number of suggestions	`5`

Returns:

Type	Description
`list[SearchResult]`	List of similar functionals

search_by_problem `async` ¶

search_by_problem(problem_description: str, domain: str | None = None, limit: int = 10) -> list[SearchResult]

Search for functionals suitable for a specific problem.

Parameters:

Name	Type	Description	Default
`problem_description`	`str`	Description of the problem to solve	required
`domain`	`str \| None`	Optional domain filter	`None`
`limit`	`int`	Maximum number of results	`10`

Returns:

Type	Description
`list[SearchResult]`	List of suitable functionals

Model Validation¶

ValidationEngine¶

Automated validation for registered neural functionals.

from opifex.platform.registry import ValidationEngine

opifex.platform.registry.validation.ValidationEngine ¶

ValidationEngine(registry_service, enable_gpu_testing: bool = True, strict_mode: bool = False)

Neural functional validation engine.

Provides full validation including functional testing, performance benchmarking, and quality assurance workflows.

Parameters:

Name	Type	Description	Default
`registry_service`		Registry service for functional access	required
`enable_gpu_testing`	`bool`	Whether to run GPU performance tests	`True`
`strict_mode`	`bool`	Whether to fail on warnings	`False`

add_rule ¶

add_rule(rule: ValidationRule) -> None

Add a validation rule.

Parameters:

Name	Type	Description	Default
`rule`	`ValidationRule`	Validation rule to add	required

remove_rule ¶

remove_rule(rule_name: str) -> bool

Remove a validation rule.

Parameters:

Name	Type	Description	Default
`rule_name`	`str`	Name of rule to remove	required

Returns:

Type	Description
`bool`	True if rule was removed

validate_functional `async` ¶

validate_functional(functional_id: str, version: str, test_types: list[TestType] | None = None, include_performance: bool = True) -> FunctionalReport

Validate a neural functional comprehensively.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of functional to validate	required
`version`	`str`	Version to validate	required
`test_types`	`list[TestType] \| None`	Types of tests to run (all if None)	`None`
`include_performance`	`bool`	Whether to include performance tests	`True`

Returns:

Type	Description
`FunctionalReport`	Complete validation report

The validation engine produces FunctionalReport objects (not ValidationReport).

Version Control¶

VersionManager¶

Git-based version control for neural functionals with dependency tracking.

from opifex.platform.registry import VersionManager

opifex.platform.registry.version.VersionManager ¶

VersionManager(storage_root: Path, registry_service, enable_git: bool = True)

Version management system for neural functionals.

Provides Git-based version control with dependency tracking, branch management, and merge capabilities.

Parameters:

Name	Type	Description	Default
`storage_root`	`Path`	Root directory for functional storage	required
`registry_service`		Registry service for data access	required
`enable_git`	`bool`	Whether to use Git for version control	`True`

create_version `async` ¶

create_version(functional_id: str, version_tag: str, author_id: str, message: str, files: dict[str, bytes], metadata: dict[str, Any] | None = None, parent_version: str | None = None, branch: str = 'main') -> Version

Create a new version of a functional.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of the functional	required
`version_tag`	`str`	Version tag (e.g., "v1.0.0")	required
`author_id`	`str`	ID of the author	required
`message`	`str`	Commit message	required
`files`	`dict[str, bytes]`	Dictionary of file names to content	required
`metadata`	`dict[str, Any] \| None`	Version metadata	`None`
`parent_version`	`str \| None`	Parent version ID	`None`
`branch`	`str`	Branch name	`'main'`

Returns:

Type	Description
`Version`	Created version object

get_version `async` ¶

get_version(functional_id: str, version_tag: str) -> Version | None

Get a specific version of a functional.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of the functional	required
`version_tag`	`str`	Version tag to retrieve	required

Returns:

Type	Description
`Version \| None`	Version object or None if not found

list_versions `async` ¶

list_versions(functional_id: str, branch: str | None = None) -> list[Version]

List all versions of a functional.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of the functional	required
`branch`	`str \| None`	Filter by branch (optional)	`None`

Returns:

Type	Description
`list[Version]`	List of versions

delete_version `async` ¶

delete_version(functional_id: str, version_tag: str, author_id: str) -> bool

Delete a version of a functional.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of the functional	required
`version_tag`	`str`	Version tag to delete	required
`author_id`	`str`	ID of the requesting user	required

Returns:

Type	Description
`bool`	True if deleted successfully

create_branch `async` ¶

create_branch(functional_id: str, branch_name: str, author_id: str, from_version: str | None = None, description: str = '') -> Branch

Create a new branch for development.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of the functional	required
`branch_name`	`str`	Name of the new branch	required
`author_id`	`str`	ID of the author	required
`from_version`	`str \| None`	Version to branch from (default: latest)	`None`
`description`	`str`	Branch description	`''`

Returns:

Type	Description
`Branch`	Created branch object

merge_branch `async` ¶

merge_branch(functional_id: str, source_branch: str, target_branch: str, author_id: str, strategy: MergeStrategy = MERGE_COMMIT, message: str = '') -> bool

Merge one branch into another.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of the functional	required
`source_branch`	`str`	Source branch to merge from	required
`target_branch`	`str`	Target branch to merge into	required
`author_id`	`str`	ID of the user performing merge	required
`strategy`	`MergeStrategy`	Merge strategy to use	`MERGE_COMMIT`
`message`	`str`	Merge commit message	`''`

Returns:

Type	Description
`bool`	True if merge was successful

get_version_diff `async` ¶

get_version_diff(functional_id: str, version_a: str, version_b: str) -> VersionDiff

Get differences between two versions.

Parameters:

Name	Type	Description	Default
`functional_id`	`str`	ID of the functional	required
`version_a`	`str`	First version	required
`version_b`	`str`	Second version	required

Returns:

Type	Description
`VersionDiff`	Differences between versions

Integration Examples¶

Complete Registry Workflow¶

import jax
from opifex.platform.registry import (
    RegistryService,
    FunctionalMetadata
)
from opifex.neural.operators.fno import FNO
from opifex.training import BasicTrainer
from opifex.data.loaders import create_darcy_loader

# Initialize registry
registry = RegistryService(storage_path="./models")

# Train model
train_loader = create_darcy_loader(
    n_samples=1000,
    batch_size=32,
    resolution=64,
    seed=42,
)
model = FNO(modes=12, width=64, depth=4)

config = TrainingConfig(num_epochs=100, learning_rate=1e-3)
trainer = BasicTrainer(model, config)
trained_model, history = trainer.train(train_loader)

# Register model
metadata = FunctionalMetadata(
    name="darcy-fno-v1",
    description="FNO for Darcy flow operator learning",
    task="operator-learning",
    domain="fluid-dynamics",
    architecture="FNO",
    tags=["pde", "elliptic", "darcy"],
    performance_metrics={
        'val_loss': history['val_loss'][-1],
        'val_relative_error': 0.015
    }
)

model_id = registry.register(model, metadata)
print(f"Model registered: {model_id}")

# Later: Load and use model
loaded_model = registry.load(model_id)
prediction = loaded_model(test_input)

# Update after fine-tuning
# ... fine-tuning code ...
new_version = registry.update(
    model_id,
    model,
    version_note="Fine-tuned on challenging cases"
)

Team Collaboration¶

# Team member 1: Register model
registry = RegistryService()
model_id = registry.register(my_model, metadata)

# Team member 2: Search and load
results = registry.search(
    query="darcy flow high accuracy",
    filters={'domain': 'fluid-dynamics'}
)
best_model_id = results[0].model_id
model = registry.load(best_model_id)

# Team member 3: Create variant
base_model = registry.load(best_model_id)
# ... modify architecture ...
variant_id = registry.register(
    modified_model,
    variant_metadata,
    parent_id=best_model_id
)

Performance and Scalability¶

Caching¶

# Enable caching for frequently accessed models
registry = RegistryService(
    enable_caching=True,
    cache_size=50  # Keep 50 models in memory
)

# Cached load (much faster for repeated access)
model = registry.load("popular-model")  # Cached after first load

Platform API Reference¶

Overview¶

Registry Core¶

RegistryService¶

Methods¶

register(model, metadata) -> str¶

load(model_id, version='latest') -> nnx.Module¶

update(model_id, model, metadata) -> str¶

search(query, filters) -> List[str]¶

list_versions(model_id) -> List[str]¶

delete(model_id, version=None)¶

Model Metadata¶

FunctionalMetadata¶

Example Usage¶

Model Search¶

SearchEngine¶

opifex.platform.registry.search.SearchEngine ¶

search async ¶

suggest_functionals async ¶

search_by_problem async ¶

Model Validation¶

ValidationEngine¶

opifex.platform.registry.validation.ValidationEngine ¶

add_rule ¶

remove_rule ¶

validate_functional async ¶

Version Control¶

VersionManager¶

opifex.platform.registry.version.VersionManager ¶

create_version async ¶

get_version async ¶

list_versions async ¶

delete_version async ¶

create_branch async ¶

merge_branch async ¶

get_version_diff async ¶

Integration Examples¶

Complete Registry Workflow¶

Team Collaboration¶

Performance and Scalability¶

Caching¶

See Also¶

`register(model, metadata) -> str`¶

`load(model_id, version='latest') -> nnx.Module`¶

`update(model_id, model, metadata) -> str`¶

`search(query, filters) -> List[str]`¶

`list_versions(model_id) -> List[str]`¶

`delete(model_id, version=None)`¶

search `async` ¶

suggest_functionals `async` ¶

search_by_problem `async` ¶

validate_functional `async` ¶

create_version `async` ¶

get_version `async` ¶

list_versions `async` ¶

delete_version `async` ¶

create_branch `async` ¶

merge_branch `async` ¶

get_version_diff `async` ¶