Wrap an Existing Model¶

Have a PyTorch STPP model already? You expose it to Seahorse by splitting it into two nn.Module adapters — a StateModel and an EventModel — that UnifiedSTPP drives. This page shows how a batch of events flows through them, what each part owns, and where an optional PresetDescriptor fits.

How a batch flows through your model¶

A padded batch of sequences enters at the top; each component's method transforms it into the next shape, ending in the per-event likelihood Seahorse trains on.

(B,T) times · (B,T,2) locations · (B,T) mask a padded batch of event sequences

StateModel.encode()

(B, T, H) hidden one hidden vector per observed event

StateModel.evolve(query_times)

(B, T_q, H) state state carried forward to each query time

EventModel.log_prob(times, locs, state, mask)

(B, T) per-event log-likelihood summed over real events → the NLL Seahorse trains and scores

Shapes: B sequences per batch, T events per sequence, T_q query times, and H = hidden_dim.

What each part owns¶

required StateModel Turns raw event history into a state vector.

encode(times, locs, mask)→ (B, T, H) hidden

evolve(hidden, query_times)→ (B, T_q, H) state

your StateModel nn.Module

required EventModel Scores events under the state — the likelihood, plus optional generation.

log_prob(times, locs, state, mask)→ (B, T) · required

sample_next(state, t_last)→ next events · optional

intensity_grid(state, t, s)→ density · optional

your EventModel nn.Module

optional PresetDescriptor Injects data-derived setup before the model is built — bounding box, coordinate stats, device fallback.

data_init_overrides(dm)→ dict merged into the build

runs once before build_model(), given the fitted data module

UnifiedSTPP(state_model, event_model, *, hidden_dim) is the wiring that ties the two modules together; the full API lives in seahorse/models/unified_model.py.

Minimal StateModel Adapter¶

import torch
import torch.nn as nn

class MyStateModel(nn.Module):
    def __init__(self, hidden_dim: int):
        super().__init__()
        self.encoder = MyExistingEncoder(hidden_dim)
        self.hidden_dim = hidden_dim

    def encode(self, times, locations, mask):
        """Encode event history.

        Args:
            times:     (B, T) event times
            locations: (B, T, 2) event locations
            mask:      (B, T) bool mask — True where events exist

        Returns:
            hidden: (B, T, hidden_dim) per-event hidden states
        """
        return self.encoder(times, locations, mask)

    def evolve(self, hidden, query_times):
        """Evolve state to query times (for piecewise-constant families: no-op).

        Args:
            hidden:      (B, T, hidden_dim)
            query_times: (B, T_q) query times

        Returns:
            state: (B, T_q, hidden_dim)
        """
        return hidden  # piecewise-constant: return last hidden state before each query

Minimal EventModel Adapter¶

class MyEventModel(nn.Module):
    def __init__(self, hidden_dim: int, spatial_dim: int = 2):
        super().__init__()
        self.decoder = MyExistingDecoder(hidden_dim, spatial_dim)

    def log_prob(self, times, locations, state, mask):
        """Compute per-event log-probability.

        Args:
            times:     (B, T) event times
            locations: (B, T, 2) event locations
            state:     (B, T, hidden_dim) evolved state at event times
            mask:      (B, T) bool mask

        Returns:
            log_prob: (B, T) per-event log-likelihood (masked positions can be 0)
        """
        return self.decoder.log_prob(times, locations, state, mask)

    def sample_next(self, state, t_last, n_samples: int = 1):
        """Optional: sample next event given state.

        Raise NotImplementedError if sampling is not supported.
        """
        raise NotImplementedError("MyEventModel does not support next-event sampling")

Wire Into a Preset¶

Once you have MyStateModel and MyEventModel, create a ModelFamilyConfig:

from dataclasses import dataclass
from seahorse.models.configs.base import BaseModelConfig, ConfigRegistry
from seahorse.models.unified_model import UnifiedSTPP

@ConfigRegistry.register("my_preset")
@dataclass
class MyPresetConfig(BaseModelConfig):
    hidden_dim: int = 64

    @classmethod
    def from_dict(cls, d, *, hidden_dim, **kwargs):
        return cls(hidden_dim=hidden_dim)

    def build_model(self) -> UnifiedSTPP:
        state = MyStateModel(self.hidden_dim)
        event = MyEventModel(self.hidden_dim)
        return UnifiedSTPP(state, event, hidden_dim=self.hidden_dim)

Then follow the Register a Preset page to expose it through the CLI and Python API.

Common Pitfalls¶

Shape mismatches: Seahorse passes (B, T, *) tensors. Check your existing model's expected input shape.
inference_mode conflict: if your model uses torch.autograd.grad internally, ensure the runner is configured with inference_mode=False.
Missing mask handling: padding positions in a batch have mask=False. Sum or mean log-prob over mask=True positions only.