Skip to main content

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::prelude::basic::{
7    RngConfig, SquareLattice, SquareLatticeConfig, SquareLatticeConfigError,
8    SquareLatticeInitMethod,
9};
10use scientific_workflow::prelude::basics::{
11    ArtifactDescriptor, ArtifactDisposition, ArtifactError, ArtifactLoadError, ExecutionScope,
12    RngRecord, RngRecordError, load_verified_artifact, persist_artifact,
13};
14use serde::{Deserialize, Serialize};
15use serde_json::{Map, Value};
16use thiserror::Error;
17
18pub const INITIALIZATION_RNG_NAMESPACE: &str = "ecological_model_core.initial_state";
19pub const INITIAL_STATE_FORMAT: &str = "ecological.initial-state.v1";
20pub const INITIAL_STATE_METADATA_KEY: &str = "initial_state";
21pub const INITIAL_STATE_ARTIFACT_REFERENCE_FORMAT: &str =
22    "ecological.initial-state-artifact-reference.v1";
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    BalancedUniform,
77    CenteredSeed,
78    CenteredDominantSeed,
79}
80
81/// Reproducible scientific instructions for one categorical state.
82#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
83#[serde(tag = "method", rename_all = "snake_case", deny_unknown_fields)]
84pub enum InitialStateRecipe {
85    Random {
86        distribution: DistributionSource,
87        #[serde(default)]
88        rng: RngConfig,
89    },
90    /// Assign the minimum-variance uniform taxon counts, then shuffle all sites.
91    BalancedUniform {
92        #[serde(default)]
93        rng: RngConfig,
94    },
95    CenteredSeed {
96        distribution: DistributionSource,
97        seed_taxon: usize,
98        seed_radius: usize,
99        #[serde(default)]
100        rng: RngConfig,
101    },
102    CenteredDominantSeed {
103        distribution: DistributionSource,
104        seed_radius: usize,
105        #[serde(default)]
106        rng: RngConfig,
107    },
108}
109
110impl InitialStateRecipe {
111    pub const fn method(&self) -> InitializationMethod {
112        match self {
113            Self::Random { .. } => InitializationMethod::Random,
114            Self::BalancedUniform { .. } => InitializationMethod::BalancedUniform,
115            Self::CenteredSeed { .. } => InitializationMethod::CenteredSeed,
116            Self::CenteredDominantSeed { .. } => InitializationMethod::CenteredDominantSeed,
117        }
118    }
119
120    pub fn distribution_path(&self) -> Option<&Path> {
121        match self {
122            Self::Random { distribution, .. }
123            | Self::CenteredSeed { distribution, .. }
124            | Self::CenteredDominantSeed { distribution, .. } => distribution.path(),
125            Self::BalancedUniform { .. } => None,
126        }
127    }
128
129    pub fn validate(
130        &self,
131        lattice: &SquareLatticeConfig,
132        num_taxa: usize,
133    ) -> Result<(), InitialStateError> {
134        if num_taxa == 0 {
135            return Err(InitialStateError::ZeroTaxa);
136        }
137        match self {
138            Self::Random { distribution, .. } => distribution.validate(num_taxa),
139            Self::BalancedUniform { .. } => Ok(()),
140            Self::CenteredSeed {
141                distribution,
142                seed_taxon,
143                seed_radius,
144                ..
145            } => {
146                distribution.validate(num_taxa)?;
147                validate_seed(lattice, num_taxa, *seed_taxon, *seed_radius)
148            }
149            Self::CenteredDominantSeed {
150                distribution,
151                seed_radius,
152                ..
153            } => {
154                distribution.validate(num_taxa)?;
155                validate_seed_geometry(lattice, *seed_radius)
156            }
157        }
158    }
159
160    pub fn create(
161        self,
162        lattice: SquareLatticeConfig,
163        num_taxa: usize,
164    ) -> Result<InitialState, InitialStateError> {
165        self.validate(&lattice, num_taxa)?;
166        let method = self.method();
167        let (distribution, rng, explicit_seed, seed_radius, dominant) = match self {
168            Self::Random { distribution, rng } => (distribution, rng, None, None, false),
169            Self::CenteredSeed {
170                distribution,
171                seed_taxon,
172                seed_radius,
173                rng,
174            } => (
175                distribution,
176                rng,
177                Some(seed_taxon),
178                Some(seed_radius),
179                false,
180            ),
181            Self::CenteredDominantSeed {
182                distribution,
183                seed_radius,
184                rng,
185            } => (distribution, rng, None, Some(seed_radius), true),
186            Self::BalancedUniform { rng } => {
187                let values = balanced_uniform_values(lattice.num_sites(), num_taxa);
188                let space = CategoricalSpace::new(
189                    lattice,
190                    SquareLatticeInitMethod::ShuffledValues { values, rng },
191                )?;
192                let counts = count_taxa(&space, num_taxa)?;
193                let rng_record = Some(rng_record_from_space(&space)?);
194                return Ok(InitialState {
195                    num_taxa,
196                    method,
197                    seed_taxon: None,
198                    rng_record,
199                    space,
200                    counts,
201                });
202            }
203        };
204        let mut weights = distribution.resolve(num_taxa)?;
205        let seed_taxon = if dominant {
206            let background = weights.get_or_insert_with(|| vec![1.0; num_taxa]);
207            Some(remove_dominant_taxon(background)?)
208        } else {
209            explicit_seed
210        };
211        let mut space = CategoricalSpace::new(
212            lattice,
213            SquareLatticeInitMethod::RandomChoices {
214                choices: (0..num_taxa).collect(),
215                weights,
216                rng,
217            },
218        )?;
219        let mut counts = count_taxa(&space, num_taxa)?;
220        if let (Some(taxon), Some(radius)) = (seed_taxon, seed_radius) {
221            let shape = space.config().shape().to_vec();
222            plant_centered_seed(&mut space, &mut counts, &shape, taxon, radius);
223        }
224        let rng_record = Some(rng_record_from_space(&space)?);
225        Ok(InitialState {
226            num_taxa,
227            method,
228            seed_taxon,
229            rng_record,
230            space,
231            counts,
232        })
233    }
234}
235
236/// A generated recipe or a verified prior-execution artifact.
237#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
238#[serde(tag = "source", rename_all = "snake_case", deny_unknown_fields)]
239pub enum InitialStateSource {
240    Recipe {
241        recipe: InitialStateRecipe,
242    },
243    VerifiedArtifact {
244        execution_directory: PathBuf,
245        descriptor: InitialStateArtifactDescriptor,
246    },
247}
248
249impl InitialStateSource {
250    pub fn validate(
251        &self,
252        lattice: &SquareLatticeConfig,
253        num_taxa: usize,
254    ) -> Result<(), InitialStateError> {
255        match self {
256            Self::Recipe { recipe } => recipe.validate(lattice, num_taxa),
257            Self::VerifiedArtifact {
258                execution_directory,
259                descriptor,
260            } => {
261                if execution_directory.as_os_str().is_empty() {
262                    return Err(InitialStateError::EmptyPath);
263                }
264                if descriptor.lattice() != lattice {
265                    return Err(InitialStateError::LatticeMismatch);
266                }
267                if descriptor.num_taxa() != num_taxa {
268                    return Err(InitialStateError::TaxonDimensionMismatch {
269                        expected: num_taxa,
270                        actual: descriptor.num_taxa(),
271                    });
272                }
273                Ok(())
274            }
275        }
276    }
277
278    pub fn resolve(
279        &self,
280        lattice: SquareLatticeConfig,
281        num_taxa: usize,
282    ) -> Result<InitialState, InitialStateError> {
283        self.validate(&lattice, num_taxa)?;
284        match self {
285            Self::Recipe { recipe } => recipe.clone().create(lattice, num_taxa),
286            Self::VerifiedArtifact {
287                execution_directory,
288                descriptor,
289            } => load_verified_initial_state(execution_directory, descriptor),
290        }
291    }
292}
293
294#[derive(Debug)]
295pub struct InitialState {
296    num_taxa: usize,
297    method: InitializationMethod,
298    seed_taxon: Option<usize>,
299    rng_record: Option<RngRecord>,
300    space: CategoricalSpace,
301    counts: TaxonCounts,
302}
303
304impl InitialState {
305    pub const fn num_taxa(&self) -> usize {
306        self.num_taxa
307    }
308    pub const fn method(&self) -> InitializationMethod {
309        self.method
310    }
311    pub const fn seed_taxon(&self) -> Option<usize> {
312        self.seed_taxon
313    }
314    pub const fn rng_record(&self) -> Option<&RngRecord> {
315        self.rng_record.as_ref()
316    }
317    pub const fn space(&self) -> &CategoricalSpace {
318        &self.space
319    }
320    pub fn counts(&self) -> &[usize] {
321        &self.counts
322    }
323    /// Returns the exact aggregate relative frequencies represented by the lattice.
324    pub fn frequencies(&self) -> Vec<f64> {
325        let total = self.space.num_sites() as f64;
326        self.counts
327            .iter()
328            .map(|&count| count as f64 / total)
329            .collect()
330    }
331    /// Encodes the complete reproducible state in eco_core's canonical JSON format.
332    pub fn to_json_bytes(&self) -> Result<Vec<u8>, serde_json::Error> {
333        serde_json::to_vec(&InitialStateDocumentRef::from(self))
334    }
335    pub fn clone_space(&self) -> CategoricalSpace {
336        self.space.clone()
337    }
338    pub fn into_parts(self) -> (CategoricalSpace, TaxonCounts) {
339        (self.space, self.counts)
340    }
341}
342
343#[derive(Serialize)]
344struct InitialStateDocumentRef<'a> {
345    format: &'static str,
346    num_taxa: usize,
347    method: InitializationMethod,
348    seed_taxon: Option<usize>,
349    rng_record: Option<&'a RngRecord>,
350    lattice: &'a SquareLatticeConfig,
351    sites: &'a [usize],
352}
353
354impl<'a> From<&'a InitialState> for InitialStateDocumentRef<'a> {
355    fn from(initial: &'a InitialState) -> Self {
356        Self {
357            format: INITIAL_STATE_FORMAT,
358            num_taxa: initial.num_taxa,
359            method: initial.method,
360            seed_taxon: initial.seed_taxon,
361            rng_record: initial.rng_record.as_ref(),
362            lattice: initial.space.config(),
363            sites: initial.space.data(),
364        }
365    }
366}
367
368#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
369#[serde(deny_unknown_fields)]
370pub struct InitialStateArtifactDescriptor {
371    format: String,
372    num_taxa: usize,
373    lattice: SquareLatticeConfig,
374    method: InitializationMethod,
375    #[serde(skip_serializing_if = "Option::is_none")]
376    seed_taxon: Option<usize>,
377    #[serde(flatten)]
378    artifact: ArtifactDescriptor,
379}
380
381/// Portable pointer to a verified initial-state artifact produced by an earlier execution.
382#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
383#[serde(deny_unknown_fields)]
384pub struct InitialStateArtifactReference {
385    format: String,
386    execution_directory: PathBuf,
387    descriptor: InitialStateArtifactDescriptor,
388}
389
390impl InitialStateArtifactReference {
391    pub fn new(
392        execution_directory: impl Into<PathBuf>,
393        descriptor: InitialStateArtifactDescriptor,
394    ) -> Self {
395        Self {
396            format: INITIAL_STATE_ARTIFACT_REFERENCE_FORMAT.to_owned(),
397            execution_directory: execution_directory.into(),
398            descriptor,
399        }
400    }
401
402    pub fn from_json_bytes(bytes: &[u8]) -> Result<Self, InitialStateError> {
403        let reference: Self =
404            serde_json::from_slice(bytes).map_err(|source| InitialStateError::Json {
405                path: PathBuf::from("<initial-state-artifact-reference>"),
406                source,
407            })?;
408        if reference.format != INITIAL_STATE_ARTIFACT_REFERENCE_FORMAT
409            || reference.execution_directory.as_os_str().is_empty()
410        {
411            return Err(InitialStateError::InvalidArtifactReference);
412        }
413        Ok(reference)
414    }
415
416    pub fn load_json(path: impl Into<PathBuf>) -> Result<Self, InitialStateError> {
417        let path = path.into();
418        let bytes = fs::read(&path).map_err(|source| InitialStateError::Io {
419            operation: "read initial-state artifact reference",
420            path: path.clone(),
421            source,
422        })?;
423        let reference: Self =
424            serde_json::from_slice(&bytes).map_err(|source| InitialStateError::Json {
425                path: path.clone(),
426                source,
427            })?;
428        if reference.format != INITIAL_STATE_ARTIFACT_REFERENCE_FORMAT
429            || reference.execution_directory.as_os_str().is_empty()
430        {
431            return Err(InitialStateError::InvalidArtifactReference);
432        }
433        Ok(reference)
434    }
435
436    pub fn to_json_bytes(&self) -> Result<Vec<u8>, serde_json::Error> {
437        serde_json::to_vec_pretty(self)
438    }
439
440    pub fn resolve(&self) -> Result<InitialState, InitialStateError> {
441        load_verified_initial_state(&self.execution_directory, &self.descriptor)
442    }
443
444    pub fn execution_directory(&self) -> &Path {
445        &self.execution_directory
446    }
447
448    pub const fn descriptor(&self) -> &InitialStateArtifactDescriptor {
449        &self.descriptor
450    }
451}
452
453impl InitialStateArtifactDescriptor {
454    pub fn format(&self) -> &str {
455        &self.format
456    }
457    pub const fn num_taxa(&self) -> usize {
458        self.num_taxa
459    }
460    pub const fn lattice(&self) -> &SquareLatticeConfig {
461        &self.lattice
462    }
463    pub const fn method(&self) -> InitializationMethod {
464        self.method
465    }
466    pub const fn seed_taxon(&self) -> Option<usize> {
467        self.seed_taxon
468    }
469    pub fn sha256(&self) -> &str {
470        self.artifact.sha256()
471    }
472    pub fn path(&self) -> &str {
473        self.artifact.path()
474    }
475    pub fn insert_into_metadata(&self, metadata: &mut Map<String, Value>) -> Option<Value> {
476        metadata.insert(
477            INITIAL_STATE_METADATA_KEY.to_owned(),
478            serde_json::to_value(self).expect("initial-state descriptor is JSON-compatible"),
479        )
480    }
481}
482
483#[derive(Clone, Debug)]
484pub struct PersistedInitialState {
485    descriptor: InitialStateArtifactDescriptor,
486    disposition: ArtifactDisposition,
487}
488
489impl PersistedInitialState {
490    pub const fn descriptor(&self) -> &InitialStateArtifactDescriptor {
491        &self.descriptor
492    }
493    pub const fn disposition(&self) -> ArtifactDisposition {
494        self.disposition
495    }
496    pub fn into_descriptor(self) -> InitialStateArtifactDescriptor {
497        self.descriptor
498    }
499}
500
501pub fn persist_initial_state(
502    scope: &ExecutionScope,
503    initial: &InitialState,
504) -> Result<PersistedInitialState, InitialStateError> {
505    let bytes = serde_json::to_vec(&InitialStateDocumentRef::from(initial))?;
506    let persisted = persist_artifact(scope, "initial-state", "json", &bytes)?;
507    Ok(PersistedInitialState {
508        descriptor: InitialStateArtifactDescriptor {
509            format: INITIAL_STATE_FORMAT.to_owned(),
510            num_taxa: initial.num_taxa,
511            lattice: initial.space.config().clone(),
512            method: initial.method,
513            seed_taxon: initial.seed_taxon,
514            artifact: persisted.descriptor().clone(),
515        },
516        disposition: persisted.disposition(),
517    })
518}
519
520pub fn load_verified_initial_state(
521    execution_directory: impl AsRef<Path>,
522    descriptor: &InitialStateArtifactDescriptor,
523) -> Result<InitialState, InitialStateError> {
524    if descriptor.format != INITIAL_STATE_FORMAT {
525        return Err(InitialStateError::UnsupportedFormat {
526            actual: descriptor.format.clone(),
527        });
528    }
529    let verified = load_verified_artifact(execution_directory, &descriptor.artifact)?;
530    let document: InitialStateDocument =
531        serde_json::from_slice(verified.bytes()).map_err(|source| InitialStateError::Json {
532            path: verified.path().to_path_buf(),
533            source,
534        })?;
535    let initial = document.resolve(descriptor.lattice.clone(), descriptor.num_taxa)?;
536    if initial.method != descriptor.method || initial.seed_taxon != descriptor.seed_taxon {
537        return Err(InitialStateError::DescriptorMismatch);
538    }
539    Ok(initial)
540}
541
542#[derive(Debug, Deserialize)]
543#[serde(deny_unknown_fields)]
544struct InitialStateDocument {
545    format: String,
546    num_taxa: usize,
547    method: InitializationMethod,
548    seed_taxon: Option<usize>,
549    rng_record: Option<RngRecord>,
550    lattice: SquareLatticeConfig,
551    sites: Vec<usize>,
552}
553
554impl InitialStateDocument {
555    fn resolve(
556        self,
557        expected_lattice: SquareLatticeConfig,
558        expected_num_taxa: usize,
559    ) -> Result<InitialState, InitialStateError> {
560        if self.format != INITIAL_STATE_FORMAT {
561            return Err(InitialStateError::UnsupportedFormat {
562                actual: self.format,
563            });
564        }
565        if self.num_taxa != expected_num_taxa {
566            return Err(InitialStateError::TaxonDimensionMismatch {
567                expected: expected_num_taxa,
568                actual: self.num_taxa,
569            });
570        }
571        if self.lattice != expected_lattice {
572            return Err(InitialStateError::LatticeMismatch);
573        }
574        if self.seed_taxon.is_some_and(|taxon| taxon >= self.num_taxa) {
575            return Err(InitialStateError::SeedTaxonOutOfRange {
576                seed_taxon: self.seed_taxon.expect("checked Some"),
577                num_taxa: self.num_taxa,
578            });
579        }
580        let space = CategoricalSpace::new(
581            self.lattice,
582            SquareLatticeInitMethod::Values { values: self.sites },
583        )?;
584        let counts = count_taxa(&space, self.num_taxa)?;
585        Ok(InitialState {
586            num_taxa: self.num_taxa,
587            method: self.method,
588            seed_taxon: self.seed_taxon,
589            rng_record: self.rng_record,
590            space,
591            counts,
592        })
593    }
594}
595
596fn rng_record_from_space(space: &CategoricalSpace) -> Result<RngRecord, InitialStateError> {
597    let config = space
598        .initialization_rng_config()
599        .ok_or(InitialStateError::MissingResolvedRngConfig)?;
600    let method = config
601        .method()
602        .ok_or(InitialStateError::MissingResolvedRngConfig)?;
603    let key = config
604        .encode_seed()
605        .ok_or(InitialStateError::MissingResolvedRngConfig)?;
606    let parameters = Map::new();
607    Ok(RngRecord::new(
608        INITIALIZATION_RNG_NAMESPACE,
609        method.name(),
610        method.version(),
611        method.seed_encoding(),
612        key,
613        Some(parameters),
614    )?)
615}
616
617fn validate_distribution(weights: &[f64], num_taxa: usize) -> Result<(), InitialStateError> {
618    if weights.len() != num_taxa {
619        return Err(InitialStateError::DistributionLength {
620            expected: num_taxa,
621            actual: weights.len(),
622        });
623    }
624    let mut total = 0.0;
625    for (taxon, &weight) in weights.iter().enumerate() {
626        if !weight.is_finite() || weight < 0.0 {
627            return Err(InitialStateError::InvalidWeight { taxon, weight });
628        }
629        total += weight;
630    }
631    if !total.is_finite() || total <= 0.0 {
632        return Err(InitialStateError::NonPositiveDistribution);
633    }
634    Ok(())
635}
636
637fn remove_dominant_taxon(weights: &mut [f64]) -> Result<usize, InitialStateError> {
638    validate_distribution(weights, weights.len())?;
639    let dominant = weights
640        .iter()
641        .enumerate()
642        .max_by(|left, right| left.1.total_cmp(right.1).then(right.0.cmp(&left.0)))
643        .map(|(taxon, _)| taxon)
644        .expect("validated distribution is nonempty");
645    let background_total = weights
646        .iter()
647        .enumerate()
648        .filter_map(|(taxon, weight)| (taxon != dominant).then_some(*weight))
649        .sum::<f64>();
650    if background_total <= 0.0 {
651        return Err(InitialStateError::NoDominantSeedBackground);
652    }
653    weights[dominant] = 0.0;
654    for weight in weights {
655        *weight /= background_total;
656    }
657    Ok(dominant)
658}
659
660fn count_taxa(space: &CategoricalSpace, num_taxa: usize) -> Result<TaxonCounts, InitialStateError> {
661    let mut counts = vec![0; num_taxa];
662    for (site, &taxon) in space.data().iter().enumerate() {
663        let count = counts
664            .get_mut(taxon)
665            .ok_or(InitialStateError::SpaceTaxonOutOfRange {
666                site,
667                taxon,
668                num_taxa,
669            })?;
670        *count += 1;
671    }
672    Ok(counts)
673}
674
675fn balanced_uniform_values(num_sites: usize, num_taxa: usize) -> Vec<usize> {
676    let per_taxon = num_sites / num_taxa;
677    let remainder = num_sites % num_taxa;
678    let mut values = Vec::with_capacity(num_sites);
679    for taxon in 0..num_taxa {
680        values.extend(std::iter::repeat_n(
681            taxon,
682            per_taxon + usize::from(taxon < remainder),
683        ));
684    }
685    debug_assert_eq!(values.len(), num_sites);
686    values
687}
688
689fn validate_seed(
690    lattice: &SquareLatticeConfig,
691    num_taxa: usize,
692    seed_taxon: usize,
693    seed_radius: usize,
694) -> Result<(), InitialStateError> {
695    if seed_taxon >= num_taxa {
696        return Err(InitialStateError::SeedTaxonOutOfRange {
697            seed_taxon,
698            num_taxa,
699        });
700    }
701    validate_seed_geometry(lattice, seed_radius)
702}
703
704fn validate_seed_geometry(
705    lattice: &SquareLatticeConfig,
706    seed_radius: usize,
707) -> Result<(), InitialStateError> {
708    let width = seed_radius
709        .checked_mul(2)
710        .and_then(|diameter| diameter.checked_add(1))
711        .ok_or(InitialStateError::SeedWidthOverflow {
712            radius: seed_radius,
713        })?;
714    for (axis, &length) in lattice.shape().iter().enumerate() {
715        if width > length {
716            return Err(InitialStateError::SeedDoesNotFit {
717                axis,
718                width,
719                length,
720            });
721        }
722    }
723    Ok(())
724}
725
726fn plant_centered_seed(
727    space: &mut CategoricalSpace,
728    counts: &mut [usize],
729    shape: &[usize],
730    seed_taxon: usize,
731    radius: usize,
732) {
733    let starts = shape
734        .iter()
735        .map(|length| length / 2 - radius)
736        .collect::<Vec<_>>();
737    let widths = vec![radius * 2 + 1; shape.len()];
738    let seed_sites = widths.iter().product::<usize>();
739    let mut coordinate = vec![0isize; shape.len()];
740    for local_flat in 0..seed_sites {
741        let mut local = local_flat;
742        for axis in (0..shape.len()).rev() {
743            coordinate[axis] = (starts[axis] + local % widths[axis]) as isize;
744            local /= widths[axis];
745        }
746        let previous = *space.get(&coordinate);
747        if previous != seed_taxon {
748            counts[previous] -= 1;
749            counts[seed_taxon] += 1;
750            space.set(&coordinate, seed_taxon);
751        }
752    }
753}
754
755#[derive(Debug, Error)]
756#[non_exhaustive]
757pub enum InitialStateError {
758    #[error("num_taxa must be positive")]
759    ZeroTaxa,
760    #[error("distribution length is {actual}, expected {expected}")]
761    DistributionLength { expected: usize, actual: usize },
762    #[error("distribution weight {taxon} is invalid: {weight}")]
763    InvalidWeight { taxon: usize, weight: f64 },
764    #[error("distribution must have positive finite total weight")]
765    NonPositiveDistribution,
766    #[error("dominant-seed initialization requires positive background mass")]
767    NoDominantSeedBackground,
768    #[error("seed taxon {seed_taxon} is outside 0..{num_taxa}")]
769    SeedTaxonOutOfRange { seed_taxon: usize, num_taxa: usize },
770    #[error("seed radius {radius} is too large")]
771    SeedWidthOverflow { radius: usize },
772    #[error("seed width {width} does not fit lattice axis {axis} of length {length}")]
773    SeedDoesNotFit {
774        axis: usize,
775        width: usize,
776        length: usize,
777    },
778    #[error("path must not be empty")]
779    EmptyPath,
780    #[error("initial-state format `{actual}` is unsupported")]
781    UnsupportedFormat { actual: String },
782    #[error("initial state declares {actual} taxa, expected {expected}")]
783    TaxonDimensionMismatch { expected: usize, actual: usize },
784    #[error("initial-state lattice does not match the expected lattice")]
785    LatticeMismatch,
786    #[error("initial-state descriptor does not match its verified document")]
787    DescriptorMismatch,
788    #[error("invalid initial-state artifact reference")]
789    InvalidArtifactReference,
790    #[error("initial-state site {site} contains taxon {taxon}, outside 0..{num_taxa}")]
791    SpaceTaxonOutOfRange {
792        site: usize,
793        taxon: usize,
794        num_taxa: usize,
795    },
796    #[error("failed to {operation} at `{path}`")]
797    Io {
798        operation: &'static str,
799        path: PathBuf,
800        #[source]
801        source: std::io::Error,
802    },
803    #[error("invalid JSON initial state at `{path}`")]
804    Json {
805        path: PathBuf,
806        #[source]
807        source: serde_json::Error,
808    },
809    #[error(transparent)]
810    Serialize(#[from] serde_json::Error),
811    #[error(transparent)]
812    Lattice(#[from] SquareLatticeConfigError),
813    #[error("initialized lattice has no resolved RNG configuration")]
814    MissingResolvedRngConfig,
815    #[error(transparent)]
816    RngRecord(#[from] RngRecordError),
817    #[error(transparent)]
818    Artifact(#[from] ArtifactError),
819    #[error(transparent)]
820    ArtifactLoad(#[from] ArtifactLoadError),
821}