ecological_model_core/
initial_state.rs

1//! Reproducible categorical ecological initial states and verified artifacts.
2
3use std::fs;
4use std::path::{Path, PathBuf};
5
6use physics_in_parallel::space::prelude::{
7    RngConfig, Space as SpaceOperations, SquareLattice, SquareLatticeConfig,
8    SquareLatticeConfigError, SquareLatticeInitMethod,
9};
10use scientific_workflow::artifact::{
11    ArtifactDescriptor, ArtifactDisposition, ArtifactError, ArtifactLoadError,
12    load_verified_artifact, persist_artifact,
13};
14use scientific_workflow::execution::ExecutionScope;
15use scientific_workflow::rng_record::{RngRecord, RngRecordError};
16use serde::{Deserialize, Serialize};
17use serde_json::{Map, Value};
18use thiserror::Error;
19
20pub const INITIALIZATION_RNG_NAMESPACE: &str = "ecological_model_core.initial_state";
21pub const INITIAL_STATE_FORMAT: &str = "ecological.initial-state.v1";
22pub const INITIAL_STATE_METADATA_KEY: &str = "initial_state";
23
24pub type CategoricalSpace = SquareLattice<usize>;
25pub type TaxonCounts = Vec<usize>;
26
27/// Source of relative taxon weights for categorical placement.
28#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
29#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
30pub enum DistributionSource {
31    Uniform,
32    Inline { weights: Vec<f64> },
33    Json { path: PathBuf },
34}
35
36impl DistributionSource {
37    pub fn path(&self) -> Option<&Path> {
38        match self {
39            Self::Json { path } => Some(path),
40            Self::Uniform | Self::Inline { .. } => None,
41        }
42    }
43
44    fn validate(&self, num_taxa: usize) -> Result<(), InitialStateError> {
45        match self {
46            Self::Uniform => Ok(()),
47            Self::Inline { weights } => validate_distribution(weights, num_taxa),
48            Self::Json { path } if path.as_os_str().is_empty() => Err(InitialStateError::EmptyPath),
49            Self::Json { .. } => Ok(()),
50        }
51    }
52
53    fn resolve(self, num_taxa: usize) -> Result<Option<Vec<f64>>, InitialStateError> {
54        let weights = match self {
55            Self::Uniform => return Ok(None),
56            Self::Inline { weights } => weights,
57            Self::Json { path } => {
58                let bytes = fs::read(&path).map_err(|source| InitialStateError::Io {
59                    operation: "read distribution",
60                    path: path.clone(),
61                    source,
62                })?;
63                serde_json::from_slice(&bytes)
64                    .map_err(|source| InitialStateError::Json { path, source })?
65            }
66        };
67        validate_distribution(&weights, num_taxa)?;
68        Ok(Some(weights))
69    }
70}
71
72#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
73#[serde(rename_all = "snake_case")]
74pub enum InitializationMethod {
75    Random,
76    CenteredSeed,
77    CenteredDominantSeed,
78}
79
80/// Reproducible scientific instructions for one categorical state.
81#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
82#[serde(tag = "method", rename_all = "snake_case", deny_unknown_fields)]
83pub enum InitialStateRecipe {
84    Random {
85        distribution: DistributionSource,
86        #[serde(default)]
87        rng: RngConfig,
88    },
89    CenteredSeed {
90        distribution: DistributionSource,
91        seed_taxon: usize,
92        seed_radius: usize,
93        #[serde(default)]
94        rng: RngConfig,
95    },
96    CenteredDominantSeed {
97        distribution: DistributionSource,
98        seed_radius: usize,
99        #[serde(default)]
100        rng: RngConfig,
101    },
102}
103
104impl InitialStateRecipe {
105    pub const fn method(&self) -> InitializationMethod {
106        match self {
107            Self::Random { .. } => InitializationMethod::Random,
108            Self::CenteredSeed { .. } => InitializationMethod::CenteredSeed,
109            Self::CenteredDominantSeed { .. } => InitializationMethod::CenteredDominantSeed,
110        }
111    }
112
113    pub fn distribution_path(&self) -> Option<&Path> {
114        match self {
115            Self::Random { distribution, .. }
116            | Self::CenteredSeed { distribution, .. }
117            | Self::CenteredDominantSeed { distribution, .. } => distribution.path(),
118        }
119    }
120
121    pub fn validate(
122        &self,
123        lattice: &SquareLatticeConfig,
124        num_taxa: usize,
125    ) -> Result<(), InitialStateError> {
126        if num_taxa == 0 {
127            return Err(InitialStateError::ZeroTaxa);
128        }
129        match self {
130            Self::Random { distribution, .. } => distribution.validate(num_taxa),
131            Self::CenteredSeed {
132                distribution,
133                seed_taxon,
134                seed_radius,
135                ..
136            } => {
137                distribution.validate(num_taxa)?;
138                validate_seed(lattice, num_taxa, *seed_taxon, *seed_radius)
139            }
140            Self::CenteredDominantSeed {
141                distribution,
142                seed_radius,
143                ..
144            } => {
145                distribution.validate(num_taxa)?;
146                validate_seed_geometry(lattice, *seed_radius)
147            }
148        }
149    }
150
151    pub fn create(
152        self,
153        lattice: SquareLatticeConfig,
154        num_taxa: usize,
155    ) -> Result<InitialState, InitialStateError> {
156        self.validate(&lattice, num_taxa)?;
157        let method = self.method();
158        let (distribution, rng, explicit_seed, seed_radius, dominant) = match self {
159            Self::Random { distribution, rng } => (distribution, rng, None, None, false),
160            Self::CenteredSeed {
161                distribution,
162                seed_taxon,
163                seed_radius,
164                rng,
165            } => (
166                distribution,
167                rng,
168                Some(seed_taxon),
169                Some(seed_radius),
170                false,
171            ),
172            Self::CenteredDominantSeed {
173                distribution,
174                seed_radius,
175                rng,
176            } => (distribution, rng, None, Some(seed_radius), true),
177        };
178        let mut weights = distribution.resolve(num_taxa)?;
179        let seed_taxon = if dominant {
180            let background = weights.get_or_insert_with(|| vec![1.0; num_taxa]);
181            Some(remove_dominant_taxon(background)?)
182        } else {
183            explicit_seed
184        };
185        let mut space = CategoricalSpace::new(
186            lattice,
187            SquareLatticeInitMethod::RandomChoices {
188                choices: (0..num_taxa).collect(),
189                weights,
190                rng,
191            },
192        )?;
193        let mut counts = count_taxa(&space, num_taxa)?;
194        if let (Some(taxon), Some(radius)) = (seed_taxon, seed_radius) {
195            let shape = space.config().shape().to_vec();
196            plant_centered_seed(&mut space, &mut counts, &shape, taxon, radius);
197        }
198        let rng_record = Some(rng_record_from_space(&space)?);
199        Ok(InitialState {
200            num_taxa,
201            method,
202            seed_taxon,
203            rng_record,
204            space,
205            counts,
206        })
207    }
208}
209
210/// A generated recipe or a verified prior-execution artifact.
211#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
212#[serde(tag = "source", rename_all = "snake_case", deny_unknown_fields)]
213pub enum InitialStateSource {
214    Recipe {
215        recipe: InitialStateRecipe,
216    },
217    VerifiedArtifact {
218        execution_directory: PathBuf,
219        descriptor: InitialStateArtifactDescriptor,
220    },
221}
222
223impl InitialStateSource {
224    pub fn validate(
225        &self,
226        lattice: &SquareLatticeConfig,
227        num_taxa: usize,
228    ) -> Result<(), InitialStateError> {
229        match self {
230            Self::Recipe { recipe } => recipe.validate(lattice, num_taxa),
231            Self::VerifiedArtifact {
232                execution_directory,
233                descriptor,
234            } => {
235                if execution_directory.as_os_str().is_empty() {
236                    return Err(InitialStateError::EmptyPath);
237                }
238                if descriptor.lattice() != lattice {
239                    return Err(InitialStateError::LatticeMismatch);
240                }
241                if descriptor.num_taxa() != num_taxa {
242                    return Err(InitialStateError::TaxonDimensionMismatch {
243                        expected: num_taxa,
244                        actual: descriptor.num_taxa(),
245                    });
246                }
247                Ok(())
248            }
249        }
250    }
251
252    pub fn resolve(
253        &self,
254        lattice: SquareLatticeConfig,
255        num_taxa: usize,
256    ) -> Result<InitialState, InitialStateError> {
257        self.validate(&lattice, num_taxa)?;
258        match self {
259            Self::Recipe { recipe } => recipe.clone().create(lattice, num_taxa),
260            Self::VerifiedArtifact {
261                execution_directory,
262                descriptor,
263            } => load_verified_initial_state(execution_directory, descriptor),
264        }
265    }
266}
267
268#[derive(Debug)]
269pub struct InitialState {
270    num_taxa: usize,
271    method: InitializationMethod,
272    seed_taxon: Option<usize>,
273    rng_record: Option<RngRecord>,
274    space: CategoricalSpace,
275    counts: TaxonCounts,
276}
277
278impl InitialState {
279    pub const fn num_taxa(&self) -> usize {
280        self.num_taxa
281    }
282    pub const fn method(&self) -> InitializationMethod {
283        self.method
284    }
285    pub const fn seed_taxon(&self) -> Option<usize> {
286        self.seed_taxon
287    }
288    pub const fn rng_record(&self) -> Option<&RngRecord> {
289        self.rng_record.as_ref()
290    }
291    pub const fn space(&self) -> &CategoricalSpace {
292        &self.space
293    }
294    pub fn counts(&self) -> &[usize] {
295        &self.counts
296    }
297    pub fn clone_space(&self) -> CategoricalSpace {
298        self.space.clone()
299    }
300    pub fn into_parts(self) -> (CategoricalSpace, TaxonCounts) {
301        (self.space, self.counts)
302    }
303}
304
305#[derive(Serialize)]
306struct InitialStateDocumentRef<'a> {
307    format: &'static str,
308    num_taxa: usize,
309    method: InitializationMethod,
310    seed_taxon: Option<usize>,
311    rng_record: Option<&'a RngRecord>,
312    lattice: &'a SquareLatticeConfig,
313    sites: &'a [usize],
314}
315
316impl<'a> From<&'a InitialState> for InitialStateDocumentRef<'a> {
317    fn from(initial: &'a InitialState) -> Self {
318        Self {
319            format: INITIAL_STATE_FORMAT,
320            num_taxa: initial.num_taxa,
321            method: initial.method,
322            seed_taxon: initial.seed_taxon,
323            rng_record: initial.rng_record.as_ref(),
324            lattice: initial.space.config(),
325            sites: initial.space.data(),
326        }
327    }
328}
329
330#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
331#[serde(deny_unknown_fields)]
332pub struct InitialStateArtifactDescriptor {
333    format: String,
334    num_taxa: usize,
335    lattice: SquareLatticeConfig,
336    method: InitializationMethod,
337    #[serde(skip_serializing_if = "Option::is_none")]
338    seed_taxon: Option<usize>,
339    #[serde(flatten)]
340    artifact: ArtifactDescriptor,
341}
342
343impl InitialStateArtifactDescriptor {
344    pub fn format(&self) -> &str {
345        &self.format
346    }
347    pub const fn num_taxa(&self) -> usize {
348        self.num_taxa
349    }
350    pub const fn lattice(&self) -> &SquareLatticeConfig {
351        &self.lattice
352    }
353    pub const fn method(&self) -> InitializationMethod {
354        self.method
355    }
356    pub const fn seed_taxon(&self) -> Option<usize> {
357        self.seed_taxon
358    }
359    pub fn sha256(&self) -> &str {
360        self.artifact.sha256()
361    }
362    pub fn path(&self) -> &str {
363        self.artifact.path()
364    }
365    pub fn insert_into_metadata(&self, metadata: &mut Map<String, Value>) -> Option<Value> {
366        metadata.insert(
367            INITIAL_STATE_METADATA_KEY.to_owned(),
368            serde_json::to_value(self).expect("initial-state descriptor is JSON-compatible"),
369        )
370    }
371}
372
373#[derive(Clone, Debug)]
374pub struct PersistedInitialState {
375    descriptor: InitialStateArtifactDescriptor,
376    disposition: ArtifactDisposition,
377}
378
379impl PersistedInitialState {
380    pub const fn descriptor(&self) -> &InitialStateArtifactDescriptor {
381        &self.descriptor
382    }
383    pub const fn disposition(&self) -> ArtifactDisposition {
384        self.disposition
385    }
386    pub fn into_descriptor(self) -> InitialStateArtifactDescriptor {
387        self.descriptor
388    }
389}
390
391pub fn persist_initial_state(
392    scope: &ExecutionScope,
393    initial: &InitialState,
394) -> Result<PersistedInitialState, InitialStateError> {
395    let bytes = serde_json::to_vec(&InitialStateDocumentRef::from(initial))?;
396    let persisted = persist_artifact(scope, "initial-state", "json", &bytes)?;
397    Ok(PersistedInitialState {
398        descriptor: InitialStateArtifactDescriptor {
399            format: INITIAL_STATE_FORMAT.to_owned(),
400            num_taxa: initial.num_taxa,
401            lattice: initial.space.config().clone(),
402            method: initial.method,
403            seed_taxon: initial.seed_taxon,
404            artifact: persisted.descriptor().clone(),
405        },
406        disposition: persisted.disposition(),
407    })
408}
409
410pub fn load_verified_initial_state(
411    execution_directory: impl AsRef<Path>,
412    descriptor: &InitialStateArtifactDescriptor,
413) -> Result<InitialState, InitialStateError> {
414    if descriptor.format != INITIAL_STATE_FORMAT {
415        return Err(InitialStateError::UnsupportedFormat {
416            actual: descriptor.format.clone(),
417        });
418    }
419    let verified = load_verified_artifact(execution_directory, &descriptor.artifact)?;
420    let document: InitialStateDocument =
421        serde_json::from_slice(verified.bytes()).map_err(|source| InitialStateError::Json {
422            path: verified.path().to_path_buf(),
423            source,
424        })?;
425    let initial = document.resolve(descriptor.lattice.clone(), descriptor.num_taxa)?;
426    if initial.method != descriptor.method || initial.seed_taxon != descriptor.seed_taxon {
427        return Err(InitialStateError::DescriptorMismatch);
428    }
429    Ok(initial)
430}
431
432#[derive(Debug, Deserialize)]
433#[serde(deny_unknown_fields)]
434struct InitialStateDocument {
435    format: String,
436    num_taxa: usize,
437    method: InitializationMethod,
438    seed_taxon: Option<usize>,
439    rng_record: Option<RngRecord>,
440    lattice: SquareLatticeConfig,
441    sites: Vec<usize>,
442}
443
444impl InitialStateDocument {
445    fn resolve(
446        self,
447        expected_lattice: SquareLatticeConfig,
448        expected_num_taxa: usize,
449    ) -> Result<InitialState, InitialStateError> {
450        if self.format != INITIAL_STATE_FORMAT {
451            return Err(InitialStateError::UnsupportedFormat {
452                actual: self.format,
453            });
454        }
455        if self.num_taxa != expected_num_taxa {
456            return Err(InitialStateError::TaxonDimensionMismatch {
457                expected: expected_num_taxa,
458                actual: self.num_taxa,
459            });
460        }
461        if self.lattice != expected_lattice {
462            return Err(InitialStateError::LatticeMismatch);
463        }
464        if self.seed_taxon.is_some_and(|taxon| taxon >= self.num_taxa) {
465            return Err(InitialStateError::SeedTaxonOutOfRange {
466                seed_taxon: self.seed_taxon.expect("checked Some"),
467                num_taxa: self.num_taxa,
468            });
469        }
470        let space = CategoricalSpace::new(
471            self.lattice,
472            SquareLatticeInitMethod::Values { values: self.sites },
473        )?;
474        let counts = count_taxa(&space, self.num_taxa)?;
475        Ok(InitialState {
476            num_taxa: self.num_taxa,
477            method: self.method,
478            seed_taxon: self.seed_taxon,
479            rng_record: self.rng_record,
480            space,
481            counts,
482        })
483    }
484}
485
486fn rng_record_from_space(space: &CategoricalSpace) -> Result<RngRecord, InitialStateError> {
487    let config = space
488        .initialization_rng_config()
489        .ok_or(InitialStateError::MissingResolvedRngConfig)?;
490    let method = config
491        .method()
492        .ok_or(InitialStateError::MissingResolvedRngConfig)?;
493    let key = config
494        .encode_seed()
495        .ok_or(InitialStateError::MissingResolvedRngConfig)?;
496    let mut parameters = Map::new();
497    if let Some(streams) = config.parallel_streams() {
498        parameters.insert("parallel_streams".to_owned(), Value::from(streams.get()));
499    }
500    Ok(RngRecord::new(
501        INITIALIZATION_RNG_NAMESPACE,
502        method.name(),
503        method.version(),
504        method.seed_encoding(),
505        key,
506        Some(parameters),
507    )?)
508}
509
510fn validate_distribution(weights: &[f64], num_taxa: usize) -> Result<(), InitialStateError> {
511    if weights.len() != num_taxa {
512        return Err(InitialStateError::DistributionLength {
513            expected: num_taxa,
514            actual: weights.len(),
515        });
516    }
517    let mut total = 0.0;
518    for (taxon, &weight) in weights.iter().enumerate() {
519        if !weight.is_finite() || weight < 0.0 {
520            return Err(InitialStateError::InvalidWeight { taxon, weight });
521        }
522        total += weight;
523    }
524    if !total.is_finite() || total <= 0.0 {
525        return Err(InitialStateError::NonPositiveDistribution);
526    }
527    Ok(())
528}
529
530fn remove_dominant_taxon(weights: &mut [f64]) -> Result<usize, InitialStateError> {
531    validate_distribution(weights, weights.len())?;
532    let dominant = weights
533        .iter()
534        .enumerate()
535        .max_by(|left, right| left.1.total_cmp(right.1).then(right.0.cmp(&left.0)))
536        .map(|(taxon, _)| taxon)
537        .expect("validated distribution is nonempty");
538    let background_total = weights.iter().sum::<f64>() - weights[dominant];
539    if background_total <= 0.0 {
540        return Err(InitialStateError::NoDominantSeedBackground);
541    }
542    weights[dominant] = 0.0;
543    for weight in weights {
544        *weight /= background_total;
545    }
546    Ok(dominant)
547}
548
549fn count_taxa(space: &CategoricalSpace, num_taxa: usize) -> Result<TaxonCounts, InitialStateError> {
550    let mut counts = vec![0; num_taxa];
551    for (site, &taxon) in space.data().iter().enumerate() {
552        let count = counts
553            .get_mut(taxon)
554            .ok_or(InitialStateError::SpaceTaxonOutOfRange {
555                site,
556                taxon,
557                num_taxa,
558            })?;
559        *count += 1;
560    }
561    Ok(counts)
562}
563
564fn validate_seed(
565    lattice: &SquareLatticeConfig,
566    num_taxa: usize,
567    seed_taxon: usize,
568    seed_radius: usize,
569) -> Result<(), InitialStateError> {
570    if seed_taxon >= num_taxa {
571        return Err(InitialStateError::SeedTaxonOutOfRange {
572            seed_taxon,
573            num_taxa,
574        });
575    }
576    validate_seed_geometry(lattice, seed_radius)
577}
578
579fn validate_seed_geometry(
580    lattice: &SquareLatticeConfig,
581    seed_radius: usize,
582) -> Result<(), InitialStateError> {
583    let width = seed_radius
584        .checked_mul(2)
585        .and_then(|diameter| diameter.checked_add(1))
586        .ok_or(InitialStateError::SeedWidthOverflow {
587            radius: seed_radius,
588        })?;
589    for (axis, &length) in lattice.shape().iter().enumerate() {
590        if width > length {
591            return Err(InitialStateError::SeedDoesNotFit {
592                axis,
593                width,
594                length,
595            });
596        }
597    }
598    Ok(())
599}
600
601fn plant_centered_seed(
602    space: &mut CategoricalSpace,
603    counts: &mut [usize],
604    shape: &[usize],
605    seed_taxon: usize,
606    radius: usize,
607) {
608    let starts = shape
609        .iter()
610        .map(|length| length / 2 - radius)
611        .collect::<Vec<_>>();
612    let widths = vec![radius * 2 + 1; shape.len()];
613    let seed_sites = widths.iter().product::<usize>();
614    let mut coordinate = vec![0isize; shape.len()];
615    for local_flat in 0..seed_sites {
616        let mut local = local_flat;
617        for axis in (0..shape.len()).rev() {
618            coordinate[axis] = (starts[axis] + local % widths[axis]) as isize;
619            local /= widths[axis];
620        }
621        let previous = *space.get(&coordinate);
622        if previous != seed_taxon {
623            counts[previous] -= 1;
624            counts[seed_taxon] += 1;
625            space.set(&coordinate, seed_taxon);
626        }
627    }
628}
629
630#[derive(Debug, Error)]
631#[non_exhaustive]
632pub enum InitialStateError {
633    #[error("num_taxa must be positive")]
634    ZeroTaxa,
635    #[error("distribution length is {actual}, expected {expected}")]
636    DistributionLength { expected: usize, actual: usize },
637    #[error("distribution weight {taxon} is invalid: {weight}")]
638    InvalidWeight { taxon: usize, weight: f64 },
639    #[error("distribution must have positive finite total weight")]
640    NonPositiveDistribution,
641    #[error("dominant-seed initialization requires positive background mass")]
642    NoDominantSeedBackground,
643    #[error("seed taxon {seed_taxon} is outside 0..{num_taxa}")]
644    SeedTaxonOutOfRange { seed_taxon: usize, num_taxa: usize },
645    #[error("seed radius {radius} is too large")]
646    SeedWidthOverflow { radius: usize },
647    #[error("seed width {width} does not fit lattice axis {axis} of length {length}")]
648    SeedDoesNotFit {
649        axis: usize,
650        width: usize,
651        length: usize,
652    },
653    #[error("path must not be empty")]
654    EmptyPath,
655    #[error("initial-state format `{actual}` is unsupported")]
656    UnsupportedFormat { actual: String },
657    #[error("initial state declares {actual} taxa, expected {expected}")]
658    TaxonDimensionMismatch { expected: usize, actual: usize },
659    #[error("initial-state lattice does not match the expected lattice")]
660    LatticeMismatch,
661    #[error("initial-state descriptor does not match its verified document")]
662    DescriptorMismatch,
663    #[error("initial-state site {site} contains taxon {taxon}, outside 0..{num_taxa}")]
664    SpaceTaxonOutOfRange {
665        site: usize,
666        taxon: usize,
667        num_taxa: usize,
668    },
669    #[error("failed to {operation} at `{path}`")]
670    Io {
671        operation: &'static str,
672        path: PathBuf,
673        #[source]
674        source: std::io::Error,
675    },
676    #[error("invalid JSON initial state at `{path}`")]
677    Json {
678        path: PathBuf,
679        #[source]
680        source: serde_json::Error,
681    },
682    #[error(transparent)]
683    Serialize(#[from] serde_json::Error),
684    #[error(transparent)]
685    Lattice(#[from] SquareLatticeConfigError),
686    #[error("initialized lattice has no resolved RNG configuration")]
687    MissingResolvedRngConfig,
688    #[error(transparent)]
689    RngRecord(#[from] RngRecordError),
690    #[error(transparent)]
691    Artifact(#[from] ArtifactError),
692    #[error(transparent)]
693    ArtifactLoad(#[from] ArtifactLoadError),
694}
ecological_model_core/initial_state.rs

ecological_model_core/
initial_state.rs
