438 lines
16 KiB
Rust
438 lines
16 KiB
Rust
//! Connectome artifact format (schema 1) and the directory loader.
|
|
//!
|
|
//! Ports `dataset/format.ts` and `dataset/load-node.ts`, including the exact fingerprint string:
|
|
//! seven lowercase SHA-256 digests joined with `:`, over `JSON.stringify(meta)` (circuit roles
|
|
//! already merged) and then the six simulation arrays in their frozen order.
|
|
|
|
use std::fs;
|
|
use std::io::Read as _;
|
|
use std::path::Path;
|
|
|
|
use indexmap::IndexMap;
|
|
use sha2::{Digest, Sha256};
|
|
|
|
use crate::error::{bail, Error, Result};
|
|
use crate::json::JsonValue;
|
|
|
|
/// `meta.visual`: the retina population.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct VisualMeta {
|
|
pub population: String,
|
|
pub count: usize,
|
|
}
|
|
|
|
/// Anatomical role name -> sorted neuron indices. Labels, not inferred task functions.
|
|
pub type Roles = IndexMap<String, Vec<u32>>;
|
|
|
|
/// The typed view of `meta.json` the kernel needs.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct BrainMetadata {
|
|
pub schema_version: i64,
|
|
pub dataset: String,
|
|
pub neurons: usize,
|
|
pub edges: usize,
|
|
/// Role order is observable: it decides the network's tracked-role order.
|
|
pub roles: Roles,
|
|
pub visual: VisualMeta,
|
|
}
|
|
|
|
/// A directed, signed, aggregated connectivity graph in source-indexed CSR form, plus roles and
|
|
/// the retina columns.
|
|
#[derive(Debug, Clone)]
|
|
pub struct BrainDataset {
|
|
/// SHA-256 digests of metadata and every array, joined with ':'; used for checkpoint
|
|
/// compatibility. `None` for a dataset assembled in memory, as in TypeScript.
|
|
pub fingerprint: Option<String>,
|
|
pub meta: BrainMetadata,
|
|
/// CSR row pointers, length neurons + 1.
|
|
pub indptr: Vec<u32>,
|
|
/// CSR column indices (post-synaptic neuron per edge).
|
|
pub targets: Vec<u32>,
|
|
/// Signed aggregated synapse counts per edge (inhibitory transmitters negative).
|
|
pub weights: Vec<i16>,
|
|
/// Neuron index of each retina column.
|
|
pub visual_indices: Vec<u32>,
|
|
/// 0 = left (mirrored on X when projecting), 1 = right.
|
|
pub visual_hemisphere: Vec<u8>,
|
|
/// Interleaved x,y column coordinates in dataset units.
|
|
pub visual_xy: Vec<f32>,
|
|
}
|
|
|
|
impl BrainDataset {
|
|
/// Neuron indices of a role, or an empty slice when the dataset does not declare it.
|
|
pub fn role(&self, name: &str) -> &[u32] {
|
|
self.meta
|
|
.roles
|
|
.get(name)
|
|
.map(|neurons| neurons.as_slice())
|
|
.unwrap_or(&[])
|
|
}
|
|
}
|
|
|
|
/// Throw if the arrays do not describe the connectome the metadata claims.
|
|
pub fn validate_dataset(dataset: &BrainDataset) -> Result<()> {
|
|
let meta = &dataset.meta;
|
|
if dataset.indptr.len() != meta.neurons + 1
|
|
|| dataset.targets.len() != meta.edges
|
|
|| dataset.weights.len() != meta.edges
|
|
{
|
|
bail!("FlyWire artifact lengths do not match metadata");
|
|
}
|
|
if dataset.visual_indices.len() != meta.visual.count
|
|
|| dataset.visual_hemisphere.len() != meta.visual.count
|
|
|| dataset.visual_xy.len() != meta.visual.count * 2
|
|
{
|
|
bail!("FlyWire visual artifact lengths do not match metadata");
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Prefix of the macro-type populations (`docs/design/macros.md` section 11).
|
|
///
|
|
/// The one name rule that splits `circuit-roles.json` in two at load time: everything else is an
|
|
/// anatomical role that the dataset fingerprint covers, and a `macro_*` role is a relabelling of
|
|
/// neurons the dataset already carries, merged *after* the fingerprint is taken. See
|
|
/// [`merge_macro_roles`].
|
|
pub const MACRO_ROLE_PREFIX: &str = "macro_";
|
|
|
|
/// Merge the sidecar's anatomical circuit roles into loaded metadata, in place.
|
|
///
|
|
/// Both the fingerprint and the role order depend on this being an `Object.assign` over the
|
|
/// existing `roles` object rather than a rebuild: an existing key (`descending`, `motor`) keeps
|
|
/// its position and a new one (`sensory`, `kenyon`, `mbon`) is appended.
|
|
///
|
|
/// `macro_*` roles are deliberately *not* merged here: [`merge_macro_roles`] adds them after the
|
|
/// fingerprint has been taken.
|
|
pub fn merge_circuit_roles(meta: &mut JsonValue, circuits: &JsonValue) -> Result<()> {
|
|
let declared = meta
|
|
.get("neurons")
|
|
.and_then(JsonValue::as_usize)
|
|
.ok_or_else(|| Error::new("Unable to load brain metadata: neurons is missing"))?;
|
|
let sidecar = circuits.get("neurons").and_then(JsonValue::as_usize);
|
|
if sidecar != Some(declared) {
|
|
bail!("Circuit roles do not match connectome");
|
|
}
|
|
let incoming = circuits
|
|
.get("roles")
|
|
.and_then(JsonValue::as_object)
|
|
.map(<[(String, JsonValue)]>::to_vec)
|
|
.unwrap_or_default();
|
|
let Some(JsonValue::Object(_)) = meta.get("roles") else {
|
|
bail!("Unable to load brain metadata: roles is missing");
|
|
};
|
|
// Object.assign(meta.roles, circuits.roles), minus the macro populations
|
|
let mut roles = meta.get("roles").cloned().unwrap_or_else(JsonValue::object);
|
|
for (name, value) in incoming {
|
|
if name.starts_with(MACRO_ROLE_PREFIX) {
|
|
continue;
|
|
}
|
|
roles.set(&name, value);
|
|
}
|
|
meta.set("roles", roles);
|
|
Ok(())
|
|
}
|
|
|
|
/// Merge the sidecar's `macro_<type>` populations into loaded metadata, in place.
|
|
///
|
|
/// Split from [`merge_circuit_roles`] for one reason, and it is a contract rather than a
|
|
/// convenience: the dataset fingerprint hashes `JSON.stringify(meta)`, checkpoints record that
|
|
/// string, and the macro populations name no new neuron, edge or weight — they are the mushroom
|
|
/// body output neurons and the brain motor neurons the artifact already listed, relabelled
|
|
/// twenty-two ways (`docs/design/macros.md` section 11: "the neuron ids, edges and kernel are
|
|
/// untouched, so the compatibility string must not move"). So a loader takes the fingerprint over
|
|
/// the anatomical roles alone and calls this afterwards, which is what keeps
|
|
/// `flysim --print-compatibility` byte-identical across this change and what lets a checkpoint
|
|
/// written before the roles existed load and start them at zero.
|
|
///
|
|
/// The cost is stated rather than hidden: re-cutting the macro populations differently would not
|
|
/// invalidate a checkpoint, and the rates restored by name would then belong to different
|
|
/// neurons. `tools/build_flywire.py` owns that partition and
|
|
/// `flybrain-core/tests/macro_roles.rs` pins it against the committed artifact.
|
|
pub fn merge_macro_roles(meta: &mut JsonValue, circuits: &JsonValue) -> Result<()> {
|
|
let incoming = circuits
|
|
.get("roles")
|
|
.and_then(JsonValue::as_object)
|
|
.map(<[(String, JsonValue)]>::to_vec)
|
|
.unwrap_or_default();
|
|
let mut roles = meta.get("roles").cloned().unwrap_or_else(JsonValue::object);
|
|
for (name, value) in incoming {
|
|
if name.starts_with(MACRO_ROLE_PREFIX) {
|
|
roles.set(&name, value);
|
|
}
|
|
}
|
|
meta.set("roles", roles);
|
|
Ok(())
|
|
}
|
|
|
|
/// Read the typed metadata out of a merged `meta.json` value.
|
|
pub fn metadata_from_json(meta: &JsonValue) -> Result<BrainMetadata> {
|
|
let field = |name: &str| -> Result<&JsonValue> {
|
|
meta.get(name)
|
|
.ok_or_else(|| Error::new(format!("Unable to load brain metadata: {name} is missing")))
|
|
};
|
|
let schema_version = field("schemaVersion")?
|
|
.as_f64()
|
|
.ok_or_else(|| Error::new("Unable to load brain metadata: schemaVersion is not a number"))?
|
|
as i64;
|
|
let dataset = field("dataset")?
|
|
.as_str()
|
|
.ok_or_else(|| Error::new("Unable to load brain metadata: dataset is not a string"))?
|
|
.to_string();
|
|
let neurons = field("neurons")?
|
|
.as_usize()
|
|
.ok_or_else(|| Error::new("Unable to load brain metadata: neurons is not an integer"))?;
|
|
let edges = field("edges")?
|
|
.as_usize()
|
|
.ok_or_else(|| Error::new("Unable to load brain metadata: edges is not an integer"))?;
|
|
let mut roles = Roles::new();
|
|
for (name, value) in field("roles")?
|
|
.as_object()
|
|
.ok_or_else(|| Error::new("Unable to load brain metadata: roles is not an object"))?
|
|
{
|
|
let mut neurons = Vec::new();
|
|
for entry in value
|
|
.as_array()
|
|
.ok_or_else(|| Error::new(format!("Role {name} is not an array")))?
|
|
{
|
|
neurons.push(
|
|
entry
|
|
.as_usize()
|
|
.ok_or_else(|| Error::new(format!("Role {name} holds a non-index")))?
|
|
as u32,
|
|
);
|
|
}
|
|
roles.insert(name.clone(), neurons);
|
|
}
|
|
let visual = field("visual")?;
|
|
let visual = VisualMeta {
|
|
population: visual
|
|
.get("population")
|
|
.and_then(JsonValue::as_str)
|
|
.unwrap_or_default()
|
|
.to_string(),
|
|
count: visual
|
|
.get("count")
|
|
.and_then(JsonValue::as_usize)
|
|
.ok_or_else(|| Error::new("Unable to load brain metadata: visual.count is missing"))?,
|
|
};
|
|
Ok(BrainMetadata {
|
|
schema_version,
|
|
dataset,
|
|
neurons,
|
|
edges,
|
|
roles,
|
|
visual,
|
|
})
|
|
}
|
|
|
|
/// Lowercase hex SHA-256 of `bytes`.
|
|
pub fn sha256_hex(bytes: &[u8]) -> String {
|
|
let digest = Sha256::digest(bytes);
|
|
let mut out = String::with_capacity(64);
|
|
for byte in digest {
|
|
out.push_str(&format!("{byte:02x}"));
|
|
}
|
|
out
|
|
}
|
|
|
|
/// The metadata the fingerprint is taken over: everything, minus the `macro_*` populations.
|
|
///
|
|
/// Key order is what a digest is made of, so this rebuilds `roles` in its own place and in its own
|
|
/// order — the macro roles are appended last by the artifact, so dropping them leaves exactly the
|
|
/// object the fingerprint was defined over before they existed. The TypeScript oracle does the
|
|
/// same in `fingerprintedMetadata` (`dataset/format.ts`). See [`merge_macro_roles`].
|
|
pub fn fingerprinted_metadata(meta: &JsonValue) -> String {
|
|
let Some(JsonValue::Object(entries)) = meta.get("roles") else {
|
|
return meta.stringify();
|
|
};
|
|
if !entries
|
|
.iter()
|
|
.any(|(name, _)| name.starts_with(MACRO_ROLE_PREFIX))
|
|
{
|
|
return meta.stringify();
|
|
}
|
|
let mut roles = JsonValue::object();
|
|
for (name, value) in entries.clone() {
|
|
if !name.starts_with(MACRO_ROLE_PREFIX) {
|
|
roles.set(&name, value);
|
|
}
|
|
}
|
|
let mut trimmed = meta.clone();
|
|
trimmed.set("roles", roles);
|
|
trimmed.stringify()
|
|
}
|
|
|
|
/// Digest of the metadata JSON and every array, joined with ':'.
|
|
///
|
|
/// The order of the seven parts and the hashed byte ranges are frozen, because checkpoints record
|
|
/// this string: metadata as UTF-8 `JSON.stringify(meta)` (circuit roles already merged), then
|
|
/// indptr, targets, weights, visualIndices, visualHemisphere, visualXY, each as its raw
|
|
/// little-endian typed-array bytes.
|
|
pub fn fingerprint_dataset(meta_json: &str, dataset: &BrainDataset) -> String {
|
|
[
|
|
sha256_hex(meta_json.as_bytes()),
|
|
sha256_hex(&to_le_bytes_u32(&dataset.indptr)),
|
|
sha256_hex(&to_le_bytes_u32(&dataset.targets)),
|
|
sha256_hex(&to_le_bytes_i16(&dataset.weights)),
|
|
sha256_hex(&to_le_bytes_u32(&dataset.visual_indices)),
|
|
sha256_hex(&dataset.visual_hemisphere),
|
|
sha256_hex(&to_le_bytes_f32(&dataset.visual_xy)),
|
|
]
|
|
.join(":")
|
|
}
|
|
|
|
fn to_le_bytes_u32(values: &[u32]) -> Vec<u8> {
|
|
let mut out = Vec::with_capacity(values.len() * 4);
|
|
for value in values {
|
|
out.extend_from_slice(&value.to_le_bytes());
|
|
}
|
|
out
|
|
}
|
|
|
|
fn to_le_bytes_i16(values: &[i16]) -> Vec<u8> {
|
|
let mut out = Vec::with_capacity(values.len() * 2);
|
|
for value in values {
|
|
out.extend_from_slice(&value.to_le_bytes());
|
|
}
|
|
out
|
|
}
|
|
|
|
fn to_le_bytes_f32(values: &[f32]) -> Vec<u8> {
|
|
let mut out = Vec::with_capacity(values.len() * 4);
|
|
for value in values {
|
|
out.extend_from_slice(&value.to_le_bytes());
|
|
}
|
|
out
|
|
}
|
|
|
|
/// Load a dataset from a local `data/<dataset>` directory, exactly as `load-node.ts` does.
|
|
pub fn load_brain_dataset_from_dir(dir: impl AsRef<Path>) -> Result<BrainDataset> {
|
|
let dir = dir.as_ref();
|
|
let meta_text = fs::read_to_string(dir.join("meta.json"))
|
|
.map_err(|error| Error::new(format!("Unable to load brain metadata: {error}")))?;
|
|
let mut meta_json = JsonValue::parse(&meta_text)
|
|
.map_err(|error| Error::new(format!("Unable to load brain metadata: {error}")))?;
|
|
let circuits_text = fs::read_to_string(dir.join("circuit-roles.json"))
|
|
.map_err(|_| Error::new("Unable to load anatomical circuit roles"))?;
|
|
let circuits = JsonValue::parse(&circuits_text)
|
|
.map_err(|_| Error::new("Unable to load anatomical circuit roles"))?;
|
|
merge_circuit_roles(&mut meta_json, &circuits)?;
|
|
merge_macro_roles(&mut meta_json, &circuits)?;
|
|
let meta = metadata_from_json(&meta_json)?;
|
|
|
|
let dataset = BrainDataset {
|
|
fingerprint: None,
|
|
meta,
|
|
indptr: load_u32(dir, "indptr.binz")?,
|
|
targets: load_u32(dir, "targets.binz")?,
|
|
weights: load_i16(dir, "weights.binz")?,
|
|
visual_indices: load_u32(dir, "visual-indices.binz")?,
|
|
visual_hemisphere: load_bytes(dir, "visual-hemisphere.binz")?,
|
|
visual_xy: load_f32(dir, "visual-xy.binz")?,
|
|
};
|
|
validate_dataset(&dataset)?;
|
|
// The macro populations are deliberately outside the digest (`merge_macro_roles`).
|
|
let fingerprint = fingerprint_dataset(&fingerprinted_metadata(&meta_json), &dataset);
|
|
Ok(BrainDataset {
|
|
fingerprint: Some(fingerprint),
|
|
..dataset
|
|
})
|
|
}
|
|
|
|
/// Gunzip one `.binz` artifact.
|
|
fn load_bytes(dir: &Path, name: &str) -> Result<Vec<u8>> {
|
|
let path = dir.join(name);
|
|
let compressed = fs::read(&path).map_err(|error| {
|
|
Error::new(format!(
|
|
"Unable to load {}: {error}",
|
|
path.to_string_lossy()
|
|
))
|
|
})?;
|
|
let mut decoder = flate2::read::GzDecoder::new(compressed.as_slice());
|
|
let mut out = Vec::new();
|
|
decoder.read_to_end(&mut out).map_err(|error| {
|
|
Error::new(format!(
|
|
"Unable to load {}: {error}",
|
|
path.to_string_lossy()
|
|
))
|
|
})?;
|
|
Ok(out)
|
|
}
|
|
|
|
fn load_u32(dir: &Path, name: &str) -> Result<Vec<u32>> {
|
|
let bytes = load_bytes(dir, name)?;
|
|
if bytes.len() % 4 != 0 {
|
|
bail!("Artifact {name} has a partial element");
|
|
}
|
|
Ok(bytes
|
|
.chunks_exact(4)
|
|
.map(|chunk| u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
|
|
.collect())
|
|
}
|
|
|
|
fn load_i16(dir: &Path, name: &str) -> Result<Vec<i16>> {
|
|
let bytes = load_bytes(dir, name)?;
|
|
if bytes.len() % 2 != 0 {
|
|
bail!("Artifact {name} has a partial element");
|
|
}
|
|
Ok(bytes
|
|
.chunks_exact(2)
|
|
.map(|chunk| i16::from_le_bytes([chunk[0], chunk[1]]))
|
|
.collect())
|
|
}
|
|
|
|
fn load_f32(dir: &Path, name: &str) -> Result<Vec<f32>> {
|
|
let bytes = load_bytes(dir, name)?;
|
|
if bytes.len() % 4 != 0 {
|
|
bail!("Artifact {name} has a partial element");
|
|
}
|
|
Ok(bytes
|
|
.chunks_exact(4)
|
|
.map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
|
|
.collect())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn sha256_matches_the_reference() {
|
|
assert_eq!(
|
|
sha256_hex(b""),
|
|
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
|
|
);
|
|
assert_eq!(
|
|
sha256_hex(b"abc"),
|
|
"ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn merge_appends_new_roles_and_keeps_existing_positions() {
|
|
let mut meta = JsonValue::parse(
|
|
r#"{"neurons":4,"roles":{"motor":[3],"command_0":[3]},"visual":{"count":0}}"#,
|
|
)
|
|
.unwrap();
|
|
let circuits =
|
|
JsonValue::parse(r#"{"neurons":4,"roles":{"motor":[2],"kenyon":[0]}}"#).unwrap();
|
|
merge_circuit_roles(&mut meta, &circuits).unwrap();
|
|
assert_eq!(
|
|
meta.get("roles").unwrap().stringify(),
|
|
r#"{"motor":[2],"command_0":[3],"kenyon":[0]}"#
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn merge_rejects_a_neuron_count_mismatch() {
|
|
let mut meta =
|
|
JsonValue::parse(r#"{"neurons":4,"roles":{},"visual":{"count":0}}"#).unwrap();
|
|
let circuits = JsonValue::parse(r#"{"neurons":5,"roles":{}}"#).unwrap();
|
|
assert_eq!(
|
|
merge_circuit_roles(&mut meta, &circuits)
|
|
.unwrap_err()
|
|
.message(),
|
|
"Circuit roles do not match connectome"
|
|
);
|
|
}
|
|
}
|