braintools.file module

File I/O and Checkpointing Utilities.

This module provides utilities for file input/output operations and model checkpointing, specifically designed for neuroscience and machine learning workflows.

Key Features:

• MATLAB File I/O: Load and parse .mat files with automatic type conversion

• Model Checkpointing: Save and restore model states using efficient msgpack serialization

• Async Saving: Non-blocking checkpoint saves for better performance

• Type Safety: Proper handling of BrainUnit quantities and BrainState objects

• Flexible Loading: Support for mismatch handling between saved and current model structures

Quick Start:

from braintools.file import load_matfile, msgpack_save, msgpack_load

# Load MATLAB data
data = load_matfile('experiment_data.mat')

# Save model checkpoint
msgpack_save('model_checkpoint.msgpack', model_state)

# Load checkpoint back
restored_state = msgpack_load('model_checkpoint.msgpack', target=model_state)

MATLAB File Loading:

from braintools.file import load_matfile, save_matfile

# Load with default settings (excludes MATLAB headers)
data = load_matfile('data.mat')

# Include MATLAB metadata
data = load_matfile('data.mat', include_header=True)

# Access nested structures (automatically converted to Python dicts/lists)
spike_times = data['trial_data']['spike_times']

# Save data back to a MATLAB .mat file
save_matfile('out.mat', {'spike_times': spike_times})

Model Checkpointing:

import brainstate as bst
from braintools.file import msgpack_save, msgpack_load, AsyncManager

# Simple synchronous save
msgpack_save('checkpoint.msgpack', model.state_dict())

# Load checkpoint with mismatch handling
state = msgpack_load('checkpoint.msgpack', target=model.state_dict(), mismatch='warn')

# Async saving for large models (non-blocking)
with AsyncManager() as manager:
    msgpack_save('checkpoint.msgpack', model.state_dict(), async_manager=manager)
    # Continue training while save happens in background

# Custom serialization for user-defined types
from braintools.file import msgpack_register_serialization

def my_type_to_dict(obj):
    return {'data': obj.data}

def my_type_from_dict(obj, state_dict, mismatch='error'):
    obj.data = state_dict['data']
    return obj

msgpack_register_serialization(MyType, my_type_to_dict, my_type_from_dict)

Utilities for loading and saving experiment artifacts, including MATLAB archives and high-performance MsgPack checkpoints.

MATLAB I/O

load_matfile() reads a .mat file (via scipy.io.loadmat()), recursively converting MATLAB structs to dicts and cell arrays to lists. Pass include_header=True to keep the __header__ / __version__ / __globals__ metadata. MATLAB v7.3 (HDF5) files are not supported and raise NotImplementedError. save_matfile() writes a dict of variables back to a .mat file.

Checkpointing

msgpack_save() / msgpack_load() serialize PyTrees (including brainunit.Quantity and brainstate.State leaves) to and from msgpack. Notes:

• Mismatch handling. When a target is given, mismatch controls what happens on a structural difference (dict keys, list/tuple length, namedtuple fields, unit, array shape): 'error' (default) raises, 'warn' warns and keeps the target’s value, 'ignore' keeps it silently.

• In-place State restore. State leaves are restored in place: the template’s .value is mutated. Pass a throwaway template to preserve the original.

• Large arrays. Arrays above ~1 GiB are transparently chunked to bypass the msgpack 2 GiB leaf limit. msgpack_load accepts an optional max_size guard (None = unlimited).

• Async saves. AsyncManager runs the file write in the background; serialization happens synchronously so the on-disk snapshot is consistent.

• Custom types. Register handlers with msgpack_register_serialization().

load_matfile	A simple function to load a .mat file using scipy from Python.
save_matfile	Save a dictionary of variables to a MATLAB .mat file via scipy.
msgpack_from_state_dict	Restores the state of the given target using a state dict.
msgpack_to_state_dict	Returns a dictionary with the state of the given target.
msgpack_register_serialization	Register a type for serialization.
msgpack_save	Save a checkpoint of the model.
msgpack_load	Load the checkpoint from the given path using the msgpack library.
AsyncManager	A simple object to track async checkpointing.
AlreadyExistsError	Raised when a rename would overwrite an existing file.
InvalidCheckpointPath	Raised for an invalid checkpoint target or corrupt checkpoint data.