Save, Restore, and Export a Fitted Model¶

Persistence is useful when fitting is expensive, when evaluation happens in a separate pipeline, or when an analysis needs an exact record of the selected regularization and spectral settings.

In [1]:

import matplotlib
matplotlib.use("module://matplotlib_inline.backend_inline")
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import fftconvolve

from fftrf import TRF, pearsonr, r2_score

import tempfile
from pathlib import Path

import matplotlib matplotlib.use("module://matplotlib_inline.backend_inline") import matplotlib.pyplot as plt import numpy as np from scipy.signal import fftconvolve from fftrf import TRF, pearsonr, r2_score import tempfile from pathlib import Path

In [2]:

def simulate_trials(*, seed=7, n_trials=6, n_samples=2048, noise_scale=0.12):
    rng = np.random.default_rng(seed)
    fs = 128.0
    tmin, tmax = 0.0, 0.25
    times = np.arange(round(tmin * fs), round(tmax * fs)) / fs
    true_kernel = (
        0.9 * np.exp(-0.5 * ((times - 0.055) / 0.014) ** 2)
        - 0.5 * np.exp(-0.5 * ((times - 0.125) / 0.024) ** 2)
    )

    stimulus, response = [], []
    for _ in range(n_trials):
        x = rng.standard_normal((n_samples, 1))
        y = fftconvolve(x[:, 0], true_kernel, mode="full")[:n_samples]
        y += noise_scale * rng.standard_normal(n_samples)
        stimulus.append(x)
        response.append(y[:, np.newaxis])
    return fs, tmin, tmax, times, true_kernel, stimulus, response

def simulate_trials(*, seed=7, n_trials=6, n_samples=2048, noise_scale=0.12): rng = np.random.default_rng(seed) fs = 128.0 tmin, tmax = 0.0, 0.25 times = np.arange(round(tmin * fs), round(tmax * fs)) / fs true_kernel = ( 0.9 * np.exp(-0.5 * ((times - 0.055) / 0.014) ** 2) - 0.5 * np.exp(-0.5 * ((times - 0.125) / 0.024) ** 2) ) stimulus, response = [], [] for _ in range(n_trials): x = rng.standard_normal((n_samples, 1)) y = fftconvolve(x[:, 0], true_kernel, mode="full")[:n_samples] y += noise_scale * rng.standard_normal(n_samples) stimulus.append(x) response.append(y[:, np.newaxis]) return fs, tmin, tmax, times, true_kernel, stimulus, response

Fit one model¶

In [3]:

fs, tmin, tmax, times, true_kernel, stimulus, response = simulate_trials()
model = TRF(direction=1)
model.train(
    stimulus[:-1],
    response[:-1],
    fs=fs,
    tmin=tmin,
    tmax=tmax,
    regularization=1e-2,
)
original_prediction = model.predict(stimulus[-1])

fs, tmin, tmax, times, true_kernel, stimulus, response = simulate_trials() model = TRF(direction=1) model.train( stimulus[:-1], response[:-1], fs=fs, tmin=tmin, tmax=tmax, regularization=1e-2, ) original_prediction = model.predict(stimulus[-1])

Save and load¶

Only load model files from a trusted source. Persistence formats that reconstruct Python objects should not be treated as safe input from untrusted users.

In [4]:

with tempfile.TemporaryDirectory() as directory:
    model_path = Path(directory) / "model.pkl"
    model.save(model_path)

    restored = TRF(direction=1)
    restored.load(model_path)
    restored_prediction = restored.predict(stimulus[-1])

print("same weights:", np.allclose(model.weights, restored.weights))
print("same prediction:", np.allclose(original_prediction, restored_prediction))
print("restored lambda:", restored.regularization)

with tempfile.TemporaryDirectory() as directory: model_path = Path(directory) / "model.pkl" model.save(model_path) restored = TRF(direction=1) restored.load(model_path) restored_prediction = restored.predict(stimulus[-1]) print("same weights:", np.allclose(model.weights, restored.weights)) print("same prediction:", np.allclose(original_prediction, restored_prediction)) print("restored lambda:", restored.regularization)

same weights: True
same prediction: True
restored lambda: 0.01

Export a shorter impulse-response window¶

to_impulse_response returns a copy over the requested lag range; it does not retrain or mutate the stored model.

In [5]:

short_weights, short_times = restored.to_impulse_response(
    tmin=0.0,
    tmax=0.12,
)
print("full kernel:", restored.weights.shape)
print("exported kernel:", short_weights.shape)

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(restored.times * 1e3, restored.weights[0, :, 0], label="Stored")
ax.plot(short_times * 1e3, short_weights[0, :, 0], linewidth=3,
        label="Exported window")
ax.set(xlabel="Lag (ms)", ylabel="Weight", title="Kernel export")
ax.legend()
plt.show()

short_weights, short_times = restored.to_impulse_response( tmin=0.0, tmax=0.12, ) print("full kernel:", restored.weights.shape) print("exported kernel:", short_weights.shape) fig, ax = plt.subplots(figsize=(8, 4)) ax.plot(restored.times * 1e3, restored.weights[0, :, 0], label="Stored") ax.plot(short_times * 1e3, short_weights[0, :, 0], linewidth=3, label="Exported window") ax.set(xlabel="Lag (ms)", ylabel="Weight", title="Kernel export") ax.legend() plt.show()

full kernel: (1, 32, 1)
exported kernel: (1, 15, 1)