1use std::collections::BTreeMap;
9
10use serde::{Deserialize, Serialize};
11
12use crate::adapter::{
13 AdapterRegistrySpec, AdapterSpec, InputPortSpec, ModelInputSpec, PlanningPolicy,
14};
15use crate::ids::{RepresentationId, SourceId, TypeId};
16use crate::model::{
17 AxisKind, AxisSpec, RepresentationSpec, SignalKind, SourceDescriptor, SourceGranularity,
18};
19use crate::plan::FitScope;
20
21pub const TYPE_DENSE_SIGNAL: &str = "dense_signal";
22pub const TYPE_TABLE: &str = "table";
23pub const TYPE_MULTI_BLOCK: &str = "multi_block";
24pub const TYPE_TIME_SERIES: &str = "time_series";
25pub const TYPE_GENOTYPE_MATRIX: &str = "genotype_matrix";
26pub const TYPE_IMAGE_RGB: &str = "image_rgb";
27pub const TYPE_GRAY_IMAGE: &str = "gray_image";
28pub const TYPE_MULTICHANNEL_IMAGE: &str = "multichannel_image";
29pub const TYPE_HYPERSPECTRAL_CUBE: &str = "hyperspectral_cube";
30pub const TYPE_LABEL_MASK: &str = "label_mask";
31pub const TYPE_METADATA: &str = "metadata";
32pub const TYPE_TARGET: &str = "target";
33pub const TYPE_MASS_SPEC: &str = "mass_spec";
34pub const TYPE_TEXT: &str = "text";
35
36pub const REPRESENTATION_SIGNAL_1D: &str = "signal_1d";
37pub const REPRESENTATION_SIGNAL_WITH_PROCESSINGS: &str = "signal_with_processings";
38pub const REPRESENTATION_RAMAN_SIGNAL: &str = "raman_signal";
39pub const REPRESENTATION_FTIR_SIGNAL: &str = "ftir_signal";
40pub const REPRESENTATION_TABULAR_NUMERIC: &str = "tabular_numeric";
41pub const REPRESENTATION_TABULAR_MIXED: &str = "tabular_mixed";
42pub const REPRESENTATION_FEATURE_BLOCK_SET: &str = "feature_block_set";
43pub const REPRESENTATION_SERIES_MV: &str = "series_mv";
44pub const REPRESENTATION_CLIMATE_SERIES_MV: &str = "climate_series_mv";
45pub const REPRESENTATION_VARIANT_MATRIX: &str = "variant_matrix";
46pub const REPRESENTATION_DOSAGE_MATRIX: &str = "dosage_matrix";
47pub const REPRESENTATION_RGB_IMAGE: &str = "rgb_image";
48pub const REPRESENTATION_GRAY_IMAGE: &str = "gray_image";
49pub const REPRESENTATION_MC_IMAGE: &str = "mc_image";
50pub const REPRESENTATION_MULTISPECTRAL_IMAGE: &str = "multispectral_image";
51pub const REPRESENTATION_CUBE_HWB: &str = "cube_hwb";
52pub const REPRESENTATION_HYPERSPECTRAL_CUBE: &str = REPRESENTATION_CUBE_HWB;
53pub const REPRESENTATION_SEGMENTATION_MASK: &str = "segmentation_mask";
54pub const REPRESENTATION_ROI_MASK: &str = "roi_mask";
55pub const REPRESENTATION_SAMPLE_METADATA: &str = "sample_metadata";
56pub const REPRESENTATION_TARGET_NUMERIC: &str = "target_numeric";
57pub const REPRESENTATION_TARGET_CATEGORICAL: &str = "target_categorical";
58pub const REPRESENTATION_TARGET_NUMERIC_MATRIX: &str = "target_numeric_matrix";
59pub const REPRESENTATION_TARGET_CATEGORICAL_MATRIX: &str = "target_categorical_matrix";
60pub const REPRESENTATION_MASS_SPECTRUM: &str = "mass_spectrum";
61pub const REPRESENTATION_TEXT_RAW: &str = "text_raw";
62pub const REPRESENTATION_TEXT_TOKEN_IDS: &str = "text_token_ids";
63
64#[derive(Clone, Copy, Debug, Eq, PartialEq)]
65#[non_exhaustive]
66pub enum BuiltinDataModel {
67 NirsSignal1d,
68 NirsSignalWithProcessings,
69 RamanSignal,
70 FtirSignal,
71 TabularNumeric,
72 TabularMixed,
73 FeatureBlockSet,
74 SeriesMultivariate,
75 ClimateSeriesMultivariate,
76 GenotypeVariantMatrix,
77 GenotypeDosageMatrix,
78 RgbImage,
79 GrayImage,
80 MultichannelImage,
81 MultispectralImage,
82 HyperspectralCube,
83 SegmentationMask,
84 RoiMask,
85 SampleMetadata,
86 TargetNumeric,
87 TargetCategorical,
88 TargetNumericMatrix,
89 TargetCategoricalMatrix,
90 MassSpectrum,
91 TextRaw,
92 TextTokenIds,
93}
94
95pub const BUILTIN_DATA_MODELS: &[BuiltinDataModel] = &[
96 BuiltinDataModel::NirsSignal1d,
97 BuiltinDataModel::NirsSignalWithProcessings,
98 BuiltinDataModel::RamanSignal,
99 BuiltinDataModel::FtirSignal,
100 BuiltinDataModel::TabularNumeric,
101 BuiltinDataModel::TabularMixed,
102 BuiltinDataModel::FeatureBlockSet,
103 BuiltinDataModel::SeriesMultivariate,
104 BuiltinDataModel::ClimateSeriesMultivariate,
105 BuiltinDataModel::GenotypeVariantMatrix,
106 BuiltinDataModel::GenotypeDosageMatrix,
107 BuiltinDataModel::RgbImage,
108 BuiltinDataModel::GrayImage,
109 BuiltinDataModel::MultichannelImage,
110 BuiltinDataModel::MultispectralImage,
111 BuiltinDataModel::HyperspectralCube,
112 BuiltinDataModel::SegmentationMask,
113 BuiltinDataModel::RoiMask,
114 BuiltinDataModel::SampleMetadata,
115 BuiltinDataModel::TargetNumeric,
116 BuiltinDataModel::TargetCategorical,
117 BuiltinDataModel::TargetNumericMatrix,
118 BuiltinDataModel::TargetCategoricalMatrix,
119 BuiltinDataModel::MassSpectrum,
120 BuiltinDataModel::TextRaw,
121 BuiltinDataModel::TextTokenIds,
122];
123
124#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
125pub struct BuiltinDataModelSpec {
126 pub key: String,
127 pub modality: String,
128 pub representation: RepresentationSpec,
129}
130
131impl BuiltinDataModelSpec {
132 pub fn source_descriptor(
133 &self,
134 source_id: impl Into<String>,
135 name: impl Into<String>,
136 sample_key: impl Into<String>,
137 granularity: SourceGranularity,
138 ) -> crate::Result<SourceDescriptor> {
139 let descriptor = SourceDescriptor {
140 id: SourceId::new(source_id)?,
141 name: name.into(),
142 type_id: self.representation.type_id.clone(),
143 modality: self.modality.clone(),
144 native_representation: self.representation.clone(),
145 sample_key: sample_key.into(),
146 granularity,
147 schema: BTreeMap::new(),
148 tags: BTreeMap::new(),
149 shape_contract: None,
150 };
151 descriptor.validate()?;
152 Ok(descriptor)
153 }
154}
155
156impl BuiltinDataModel {
157 pub fn key(self) -> &'static str {
158 match self {
159 Self::NirsSignal1d => "nirs.signal_1d",
160 Self::NirsSignalWithProcessings => "nirs.signal_with_processings",
161 Self::RamanSignal => "spectroscopy.raman_signal",
162 Self::FtirSignal => "spectroscopy.ftir_signal",
163 Self::TabularNumeric => "tabular.numeric",
164 Self::TabularMixed => "tabular.mixed",
165 Self::FeatureBlockSet => "features.block_set",
166 Self::SeriesMultivariate => "series.multivariate",
167 Self::ClimateSeriesMultivariate => "climate.series_multivariate",
168 Self::GenotypeVariantMatrix => "genotype.variant_matrix",
169 Self::GenotypeDosageMatrix => "genotype.dosage_matrix",
170 Self::RgbImage => "image.rgb",
171 Self::GrayImage => "image.gray",
172 Self::MultichannelImage => "image.multichannel",
173 Self::MultispectralImage => "image.multispectral",
174 Self::HyperspectralCube => "hyperspectral.cube",
175 Self::SegmentationMask => "image.segmentation_mask",
176 Self::RoiMask => "image.roi_mask",
177 Self::SampleMetadata => "metadata.sample",
178 Self::TargetNumeric => "target.numeric",
179 Self::TargetCategorical => "target.categorical",
180 Self::TargetNumericMatrix => "target.numeric_matrix",
181 Self::TargetCategoricalMatrix => "target.categorical_matrix",
182 Self::MassSpectrum => "mass_spec.spectrum",
183 Self::TextRaw => "text.raw",
184 Self::TextTokenIds => "text.token_ids",
185 }
186 }
187
188 pub fn modality(self) -> &'static str {
189 match self {
190 Self::NirsSignal1d | Self::NirsSignalWithProcessings => "nirs",
191 Self::RamanSignal => "raman",
192 Self::FtirSignal => "ftir",
193 Self::TabularNumeric | Self::TabularMixed => "tabular",
194 Self::FeatureBlockSet => "multi_block",
195 Self::SeriesMultivariate => "time_series",
196 Self::ClimateSeriesMultivariate => "climate",
197 Self::GenotypeVariantMatrix | Self::GenotypeDosageMatrix => "genotype",
198 Self::RgbImage
199 | Self::GrayImage
200 | Self::MultichannelImage
201 | Self::MultispectralImage => "image",
202 Self::HyperspectralCube => "hyperspectral",
203 Self::SegmentationMask | Self::RoiMask => "image_mask",
204 Self::SampleMetadata => "metadata",
205 Self::TargetNumeric
206 | Self::TargetCategorical
207 | Self::TargetNumericMatrix
208 | Self::TargetCategoricalMatrix => "target",
209 Self::MassSpectrum => "mass_spec",
210 Self::TextRaw | Self::TextTokenIds => "text",
211 }
212 }
213
214 pub fn representation(self) -> RepresentationSpec {
215 match self {
216 Self::NirsSignal1d => signal_1d(SignalKind::Unknown),
217 Self::NirsSignalWithProcessings => signal_with_processings(SignalKind::Unknown),
218 Self::RamanSignal => raman_signal(),
219 Self::FtirSignal => ftir_signal(),
220 Self::TabularNumeric => tabular_numeric(),
221 Self::TabularMixed => tabular_mixed(),
222 Self::FeatureBlockSet => feature_block_set(),
223 Self::SeriesMultivariate => series_mv(REPRESENTATION_SERIES_MV),
224 Self::ClimateSeriesMultivariate => series_mv(REPRESENTATION_CLIMATE_SERIES_MV),
225 Self::GenotypeVariantMatrix => genotype_variant_matrix(),
226 Self::GenotypeDosageMatrix => genotype_dosage_matrix(),
227 Self::RgbImage => rgb_image(),
228 Self::GrayImage => gray_image(),
229 Self::MultichannelImage => multichannel_image(),
230 Self::MultispectralImage => multispectral_image(),
231 Self::HyperspectralCube => hyperspectral_cube(),
232 Self::SegmentationMask => segmentation_mask(),
233 Self::RoiMask => roi_mask(),
234 Self::SampleMetadata => sample_metadata(),
235 Self::TargetNumeric => target_numeric(),
236 Self::TargetCategorical => target_categorical(),
237 Self::TargetNumericMatrix => target_numeric_matrix(),
238 Self::TargetCategoricalMatrix => target_categorical_matrix(),
239 Self::MassSpectrum => mass_spectrum(),
240 Self::TextRaw => text_raw(),
241 Self::TextTokenIds => text_token_ids(),
242 }
243 }
244
245 pub fn spec(self) -> BuiltinDataModelSpec {
246 BuiltinDataModelSpec {
247 key: self.key().to_string(),
248 modality: self.modality().to_string(),
249 representation: self.representation(),
250 }
251 }
252}
253
254pub fn builtin_data_model_specs() -> Vec<BuiltinDataModelSpec> {
255 BUILTIN_DATA_MODELS
256 .iter()
257 .copied()
258 .map(BuiltinDataModel::spec)
259 .collect()
260}
261
262pub fn builtin_representations() -> Vec<RepresentationSpec> {
263 BUILTIN_DATA_MODELS
264 .iter()
265 .copied()
266 .map(BuiltinDataModel::representation)
267 .collect()
268}
269
270pub fn tabular_numeric_model_input_spec() -> ModelInputSpec {
271 ModelInputSpec {
272 ports: vec![InputPortSpec {
273 name: "x".to_string(),
274 accepted_representations: vec![rid(REPRESENTATION_TABULAR_NUMERIC)],
275 accepted_types: vec![tid(TYPE_TABLE)],
276 rank: Some(2),
277 multi_source: true,
278 optional: false,
279 }],
280 default_fusion: None,
281 }
282}
283
284pub fn builtin_adapter_registry_spec() -> AdapterRegistrySpec {
285 AdapterRegistrySpec {
286 adapters: vec![
287 stateless_adapter(
288 "spectra.flatten",
289 TYPE_DENSE_SIGNAL,
290 REPRESENTATION_SIGNAL_WITH_PROCESSINGS,
291 TYPE_TABLE,
292 REPRESENTATION_TABULAR_NUMERIC,
293 10,
294 false,
295 ),
296 stateless_adapter(
297 "spectra.identity_features",
298 TYPE_DENSE_SIGNAL,
299 REPRESENTATION_SIGNAL_1D,
300 TYPE_TABLE,
301 REPRESENTATION_TABULAR_NUMERIC,
302 5,
303 false,
304 ),
305 stateless_adapter(
306 "spectra.raman_flatten",
307 TYPE_DENSE_SIGNAL,
308 REPRESENTATION_RAMAN_SIGNAL,
309 TYPE_TABLE,
310 REPRESENTATION_TABULAR_NUMERIC,
311 8,
312 false,
313 ),
314 stateless_adapter(
315 "spectra.ftir_flatten",
316 TYPE_DENSE_SIGNAL,
317 REPRESENTATION_FTIR_SIGNAL,
318 TYPE_TABLE,
319 REPRESENTATION_TABULAR_NUMERIC,
320 8,
321 false,
322 ),
323 stateless_adapter(
324 "features.blocks.flatten",
325 TYPE_MULTI_BLOCK,
326 REPRESENTATION_FEATURE_BLOCK_SET,
327 TYPE_TABLE,
328 REPRESENTATION_TABULAR_NUMERIC,
329 5,
330 false,
331 ),
332 stateless_adapter(
333 "weather.aggregate",
334 TYPE_TIME_SERIES,
335 REPRESENTATION_CLIMATE_SERIES_MV,
336 TYPE_TABLE,
337 REPRESENTATION_TABULAR_NUMERIC,
338 20,
339 true,
340 ),
341 stateless_adapter(
342 "series.aggregate",
343 TYPE_TIME_SERIES,
344 REPRESENTATION_SERIES_MV,
345 TYPE_TABLE,
346 REPRESENTATION_TABULAR_NUMERIC,
347 20,
348 true,
349 ),
350 stateless_adapter(
351 "genotype.dosage",
352 TYPE_GENOTYPE_MATRIX,
353 REPRESENTATION_VARIANT_MATRIX,
354 TYPE_GENOTYPE_MATRIX,
355 REPRESENTATION_DOSAGE_MATRIX,
356 8,
357 false,
358 ),
359 stateful_adapter(
360 "genotype.pca",
361 TYPE_GENOTYPE_MATRIX,
362 REPRESENTATION_DOSAGE_MATRIX,
363 TYPE_TABLE,
364 REPRESENTATION_TABULAR_NUMERIC,
365 30,
366 true,
367 ),
368 stateful_adapter(
369 "tabular.encoder",
370 TYPE_TABLE,
371 REPRESENTATION_TABULAR_MIXED,
372 TYPE_TABLE,
373 REPRESENTATION_TABULAR_NUMERIC,
374 15,
375 false,
376 ),
377 stateless_adapter(
378 "image.channel_stats",
379 TYPE_IMAGE_RGB,
380 REPRESENTATION_RGB_IMAGE,
381 TYPE_TABLE,
382 REPRESENTATION_TABULAR_NUMERIC,
383 25,
384 true,
385 ),
386 stateless_adapter(
387 "image.gray_stats",
388 TYPE_GRAY_IMAGE,
389 REPRESENTATION_GRAY_IMAGE,
390 TYPE_TABLE,
391 REPRESENTATION_TABULAR_NUMERIC,
392 20,
393 true,
394 ),
395 stateless_adapter(
396 "image.multichannel_stats",
397 TYPE_MULTICHANNEL_IMAGE,
398 REPRESENTATION_MC_IMAGE,
399 TYPE_TABLE,
400 REPRESENTATION_TABULAR_NUMERIC,
401 30,
402 true,
403 ),
404 stateless_adapter(
405 "image.multispectral_stats",
406 TYPE_MULTICHANNEL_IMAGE,
407 REPRESENTATION_MULTISPECTRAL_IMAGE,
408 TYPE_TABLE,
409 REPRESENTATION_TABULAR_NUMERIC,
410 30,
411 true,
412 ),
413 stateful_adapter(
414 "image.embedding",
415 TYPE_IMAGE_RGB,
416 REPRESENTATION_RGB_IMAGE,
417 TYPE_TABLE,
418 REPRESENTATION_TABULAR_NUMERIC,
419 60,
420 true,
421 ),
422 stateless_adapter(
423 "image.raw_tensor_chw",
424 TYPE_IMAGE_RGB,
425 REPRESENTATION_RGB_IMAGE,
426 TYPE_MULTICHANNEL_IMAGE,
427 REPRESENTATION_MC_IMAGE,
428 5,
429 false,
430 ),
431 stateless_adapter(
432 "hsi.spatial_mean",
433 TYPE_HYPERSPECTRAL_CUBE,
434 REPRESENTATION_HYPERSPECTRAL_CUBE,
435 TYPE_DENSE_SIGNAL,
436 REPRESENTATION_SIGNAL_1D,
437 20,
438 true,
439 ),
440 stateless_adapter(
441 "hsi.flatten",
442 TYPE_HYPERSPECTRAL_CUBE,
443 REPRESENTATION_HYPERSPECTRAL_CUBE,
444 TYPE_TABLE,
445 REPRESENTATION_TABULAR_NUMERIC,
446 80,
447 false,
448 ),
449 stateless_adapter(
450 "mask.area_features",
451 TYPE_LABEL_MASK,
452 REPRESENTATION_SEGMENTATION_MASK,
453 TYPE_TABLE,
454 REPRESENTATION_TABULAR_NUMERIC,
455 15,
456 true,
457 ),
458 stateless_adapter(
459 "mask.roi_area_features",
460 TYPE_LABEL_MASK,
461 REPRESENTATION_ROI_MASK,
462 TYPE_TABLE,
463 REPRESENTATION_TABULAR_NUMERIC,
464 10,
465 true,
466 ),
467 stateful_adapter(
468 "metadata.encoder",
469 TYPE_METADATA,
470 REPRESENTATION_SAMPLE_METADATA,
471 TYPE_TABLE,
472 REPRESENTATION_TABULAR_NUMERIC,
473 15,
474 true,
475 ),
476 stateless_adapter(
477 "ms.bin_to_table",
478 TYPE_MASS_SPEC,
479 REPRESENTATION_MASS_SPECTRUM,
480 TYPE_TABLE,
481 REPRESENTATION_TABULAR_NUMERIC,
482 25,
483 true,
484 ),
485 stateful_adapter(
486 "text.embedding",
487 TYPE_TEXT,
488 REPRESENTATION_TEXT_RAW,
489 TYPE_TABLE,
490 REPRESENTATION_TABULAR_NUMERIC,
491 50,
492 true,
493 ),
494 stateless_adapter(
495 "text.bag_of_tokens",
496 TYPE_TEXT,
497 REPRESENTATION_TEXT_TOKEN_IDS,
498 TYPE_TABLE,
499 REPRESENTATION_TABULAR_NUMERIC,
500 35,
501 true,
502 ),
503 ],
504 }
505}
506
507pub fn default_builtin_planning_policy() -> PlanningPolicy {
508 PlanningPolicy {
509 allow_lossy: true,
510 allow_stateful: true,
511 allow_supervised: false,
512 require_user_choice_on_ambiguity: true,
513 max_hops: Some(4),
514 ..PlanningPolicy::default()
515 }
516}
517
518pub fn signal_1d(signal_type: SignalKind) -> RepresentationSpec {
519 representation(
520 REPRESENTATION_SIGNAL_1D,
521 TYPE_DENSE_SIGNAL,
522 Some(2),
523 vec![
524 sample_axis(),
525 axis("wavelength", AxisKind::Wavelength, Some("nm"), None, false),
526 ],
527 RepresentationStorage::new("ndarray", Some("float64"), false, false, Some(signal_type)),
528 )
529}
530
531pub fn signal_with_processings(signal_type: SignalKind) -> RepresentationSpec {
532 representation(
533 REPRESENTATION_SIGNAL_WITH_PROCESSINGS,
534 TYPE_DENSE_SIGNAL,
535 Some(3),
536 vec![
537 sample_axis(),
538 axis("processing", AxisKind::Processing, None, None, false),
539 axis("wavelength", AxisKind::Wavelength, Some("nm"), None, false),
540 ],
541 RepresentationStorage::new("ndarray", Some("float64"), false, false, Some(signal_type)),
542 )
543}
544
545pub fn raman_signal() -> RepresentationSpec {
546 spectroscopy_signal(REPRESENTATION_RAMAN_SIGNAL)
547}
548
549pub fn ftir_signal() -> RepresentationSpec {
550 spectroscopy_signal(REPRESENTATION_FTIR_SIGNAL)
551}
552
553pub fn tabular_numeric() -> RepresentationSpec {
554 representation(
555 REPRESENTATION_TABULAR_NUMERIC,
556 TYPE_TABLE,
557 Some(2),
558 vec![
559 sample_axis(),
560 axis("feature", AxisKind::Feature, None, None, false),
561 ],
562 RepresentationStorage::new("dataframe", Some("float64"), false, false, None),
563 )
564}
565
566pub fn tabular_mixed() -> RepresentationSpec {
567 representation(
568 REPRESENTATION_TABULAR_MIXED,
569 TYPE_TABLE,
570 Some(2),
571 vec![
572 sample_axis(),
573 axis("column", AxisKind::Feature, None, None, false),
574 ],
575 RepresentationStorage::new("dataframe", None, false, false, None),
576 )
577}
578
579pub fn feature_block_set() -> RepresentationSpec {
580 representation(
581 REPRESENTATION_FEATURE_BLOCK_SET,
582 TYPE_MULTI_BLOCK,
583 Some(3),
584 vec![
585 sample_axis(),
586 axis("block", AxisKind::Feature, None, None, false),
587 axis("feature", AxisKind::Feature, None, None, true),
588 ],
589 RepresentationStorage::new("feature_block_set", Some("float64"), false, true, None),
590 )
591}
592
593pub fn series_mv(representation_id: &str) -> RepresentationSpec {
594 representation(
595 representation_id,
596 TYPE_TIME_SERIES,
597 Some(3),
598 vec![
599 sample_axis(),
600 axis("time", AxisKind::Time, None, None, true),
601 axis("variable", AxisKind::Feature, None, None, false),
602 ],
603 RepresentationStorage::new("ndarray", Some("float64"), false, true, None),
604 )
605}
606
607pub fn genotype_variant_matrix() -> RepresentationSpec {
608 representation(
609 REPRESENTATION_VARIANT_MATRIX,
610 TYPE_GENOTYPE_MATRIX,
611 Some(2),
612 vec![
613 sample_axis(),
614 axis("variant", AxisKind::Variant, None, None, false),
615 ],
616 RepresentationStorage::new("ndarray", Some("int8"), false, false, None),
617 )
618}
619
620pub fn genotype_dosage_matrix() -> RepresentationSpec {
621 representation(
622 REPRESENTATION_DOSAGE_MATRIX,
623 TYPE_GENOTYPE_MATRIX,
624 Some(2),
625 vec![
626 sample_axis(),
627 axis("variant", AxisKind::Variant, None, None, false),
628 ],
629 RepresentationStorage::new("ndarray", Some("float32"), false, false, None),
630 )
631}
632
633pub fn rgb_image() -> RepresentationSpec {
634 representation(
635 REPRESENTATION_RGB_IMAGE,
636 TYPE_IMAGE_RGB,
637 Some(4),
638 vec![
639 sample_axis(),
640 axis("height", AxisKind::Height, Some("px"), None, false),
641 axis("width", AxisKind::Width, Some("px"), None, false),
642 axis("channel", AxisKind::Channel, None, Some(3), false),
643 ],
644 RepresentationStorage::new("ndarray", Some("uint8"), false, false, None),
645 )
646}
647
648pub fn gray_image() -> RepresentationSpec {
649 representation(
650 REPRESENTATION_GRAY_IMAGE,
651 TYPE_GRAY_IMAGE,
652 Some(3),
653 vec![
654 sample_axis(),
655 axis("height", AxisKind::Height, Some("px"), None, false),
656 axis("width", AxisKind::Width, Some("px"), None, false),
657 ],
658 RepresentationStorage::new("ndarray", Some("uint8"), false, false, None),
659 )
660}
661
662pub fn multichannel_image() -> RepresentationSpec {
663 representation(
664 REPRESENTATION_MC_IMAGE,
665 TYPE_MULTICHANNEL_IMAGE,
666 Some(4),
667 vec![
668 sample_axis(),
669 axis("height", AxisKind::Height, Some("px"), None, false),
670 axis("width", AxisKind::Width, Some("px"), None, false),
671 axis("channel", AxisKind::Channel, None, None, false),
672 ],
673 RepresentationStorage::new("ndarray", Some("float32"), false, false, None),
674 )
675}
676
677pub fn multispectral_image() -> RepresentationSpec {
678 representation(
679 REPRESENTATION_MULTISPECTRAL_IMAGE,
680 TYPE_MULTICHANNEL_IMAGE,
681 Some(4),
682 vec![
683 sample_axis(),
684 axis("height", AxisKind::Height, Some("px"), None, false),
685 axis("width", AxisKind::Width, Some("px"), None, false),
686 axis("band", AxisKind::Channel, None, None, false),
687 ],
688 RepresentationStorage::new("ndarray", Some("float32"), false, false, None),
689 )
690}
691
692pub fn hyperspectral_cube() -> RepresentationSpec {
693 representation(
694 REPRESENTATION_HYPERSPECTRAL_CUBE,
695 TYPE_HYPERSPECTRAL_CUBE,
696 Some(4),
697 vec![
698 sample_axis(),
699 axis("height", AxisKind::Height, Some("px"), None, false),
700 axis("width", AxisKind::Width, Some("px"), None, false),
701 axis("band", AxisKind::Wavelength, Some("nm"), None, false),
702 ],
703 RepresentationStorage::new("ndarray", Some("float32"), false, false, None),
704 )
705}
706
707pub fn segmentation_mask() -> RepresentationSpec {
708 representation(
709 REPRESENTATION_SEGMENTATION_MASK,
710 TYPE_LABEL_MASK,
711 Some(3),
712 vec![
713 sample_axis(),
714 axis("height", AxisKind::Height, Some("px"), None, false),
715 axis("width", AxisKind::Width, Some("px"), None, false),
716 ],
717 RepresentationStorage::new("ndarray", Some("int32"), false, false, None),
718 )
719}
720
721pub fn roi_mask() -> RepresentationSpec {
722 representation(
723 REPRESENTATION_ROI_MASK,
724 TYPE_LABEL_MASK,
725 Some(3),
726 vec![
727 sample_axis(),
728 axis("height", AxisKind::Height, Some("px"), None, false),
729 axis("width", AxisKind::Width, Some("px"), None, false),
730 ],
731 RepresentationStorage::new("ndarray", Some("bool"), false, false, None),
732 )
733}
734
735pub fn sample_metadata() -> RepresentationSpec {
736 representation(
737 REPRESENTATION_SAMPLE_METADATA,
738 TYPE_METADATA,
739 Some(2),
740 vec![
741 sample_axis(),
742 axis("field", AxisKind::Feature, None, None, false),
743 ],
744 RepresentationStorage::new("dataframe", None, false, false, None),
745 )
746}
747
748pub fn target_numeric() -> RepresentationSpec {
749 representation(
750 REPRESENTATION_TARGET_NUMERIC,
751 TYPE_TARGET,
752 Some(1),
753 vec![sample_axis()],
754 RepresentationStorage::new("array", Some("float64"), false, false, None),
755 )
756}
757
758pub fn target_categorical() -> RepresentationSpec {
759 representation(
760 REPRESENTATION_TARGET_CATEGORICAL,
761 TYPE_TARGET,
762 Some(1),
763 vec![sample_axis()],
764 RepresentationStorage::new("array", Some("string"), false, false, None),
765 )
766}
767
768pub fn target_numeric_matrix() -> RepresentationSpec {
769 target_matrix(REPRESENTATION_TARGET_NUMERIC_MATRIX, Some("float64"))
770}
771
772pub fn target_categorical_matrix() -> RepresentationSpec {
773 target_matrix(REPRESENTATION_TARGET_CATEGORICAL_MATRIX, Some("string"))
774}
775
776pub fn mass_spectrum() -> RepresentationSpec {
777 representation(
778 REPRESENTATION_MASS_SPECTRUM,
779 TYPE_MASS_SPEC,
780 Some(2),
781 vec![
782 sample_axis(),
783 axis("mz", AxisKind::Feature, Some("m/z"), None, true),
784 ],
785 RepresentationStorage::new("ragged_array", Some("float64"), false, true, None),
786 )
787}
788
789pub fn text_raw() -> RepresentationSpec {
790 representation(
791 REPRESENTATION_TEXT_RAW,
792 TYPE_TEXT,
793 Some(1),
794 vec![sample_axis()],
795 RepresentationStorage::new("list", Some("string"), false, true, None),
796 )
797}
798
799pub fn text_token_ids() -> RepresentationSpec {
800 representation(
801 REPRESENTATION_TEXT_TOKEN_IDS,
802 TYPE_TEXT,
803 Some(2),
804 vec![
805 sample_axis(),
806 axis("token", AxisKind::Token, None, None, true),
807 ],
808 RepresentationStorage::new("ragged_array", Some("int32"), false, true, None),
809 )
810}
811
812fn target_matrix(representation_id: &str, dtype: Option<&str>) -> RepresentationSpec {
813 representation(
814 representation_id,
815 TYPE_TARGET,
816 Some(2),
817 vec![
818 sample_axis(),
819 axis("target", AxisKind::Target, None, None, false),
820 ],
821 RepresentationStorage::new("array", dtype, false, false, None),
822 )
823}
824
825fn spectroscopy_signal(representation_id: &str) -> RepresentationSpec {
826 representation(
827 representation_id,
828 TYPE_DENSE_SIGNAL,
829 Some(2),
830 vec![
831 sample_axis(),
832 axis(
833 "wavenumber",
834 AxisKind::Wavenumber,
835 Some("cm^-1"),
836 None,
837 false,
838 ),
839 ],
840 RepresentationStorage::new(
841 "ndarray",
842 Some("float64"),
843 false,
844 false,
845 Some(SignalKind::Unknown),
846 ),
847 )
848}
849
850struct RepresentationStorage<'a> {
851 container: &'a str,
852 dtype: Option<&'a str>,
853 sparse: bool,
854 ragged: bool,
855 signal_type: Option<SignalKind>,
856}
857
858impl<'a> RepresentationStorage<'a> {
859 fn new(
860 container: &'a str,
861 dtype: Option<&'a str>,
862 sparse: bool,
863 ragged: bool,
864 signal_type: Option<SignalKind>,
865 ) -> Self {
866 Self {
867 container,
868 dtype,
869 sparse,
870 ragged,
871 signal_type,
872 }
873 }
874}
875
876fn representation(
877 id: &str,
878 type_id: &str,
879 rank: Option<usize>,
880 axes: Vec<AxisSpec>,
881 storage: RepresentationStorage<'_>,
882) -> RepresentationSpec {
883 RepresentationSpec {
884 id: rid(id),
885 type_id: tid(type_id),
886 rank,
887 axes,
888 container: storage.container.to_string(),
889 dtype: storage.dtype.map(str::to_string),
890 sparse: storage.sparse,
891 ragged: storage.ragged,
892 signal_type: storage.signal_type,
893 }
894}
895
896fn sample_axis() -> AxisSpec {
897 axis("sample", AxisKind::Sample, None, None, false)
898}
899
900fn axis(
901 name: &str,
902 kind: AxisKind,
903 unit: Option<&str>,
904 size: Option<usize>,
905 variable: bool,
906) -> AxisSpec {
907 AxisSpec {
908 name: name.to_string(),
909 kind,
910 unit: unit.map(str::to_string),
911 size,
912 variable,
913 coordinate: None,
914 }
915}
916
917fn stateless_adapter(
918 id: &str,
919 input_type: &str,
920 input_representation: &str,
921 output_type: &str,
922 output_representation: &str,
923 cost: u64,
924 lossy: bool,
925) -> AdapterSpec {
926 adapter(
927 id,
928 (input_type, input_representation),
929 (output_type, output_representation),
930 cost,
931 lossy,
932 false,
933 FitScope::Stateless,
934 )
935}
936
937fn stateful_adapter(
938 id: &str,
939 input_type: &str,
940 input_representation: &str,
941 output_type: &str,
942 output_representation: &str,
943 cost: u64,
944 lossy: bool,
945) -> AdapterSpec {
946 adapter(
947 id,
948 (input_type, input_representation),
949 (output_type, output_representation),
950 cost,
951 lossy,
952 true,
953 FitScope::FoldTrain,
954 )
955}
956
957fn adapter(
958 id: &str,
959 input: (&str, &str),
960 output: (&str, &str),
961 cost: u64,
962 lossy: bool,
963 stateful: bool,
964 fit_scope: FitScope,
965) -> AdapterSpec {
966 AdapterSpec {
967 id: id.to_string(),
968 version: "1.0.0".to_string(),
969 input_type: tid(input.0),
970 input_representation: rid(input.1),
971 output_type: tid(output.0),
972 output_representation: rid(output.1),
973 cost,
974 lossy,
975 supervised: false,
976 stateful,
977 deterministic: true,
978 fit_scope,
979 params: BTreeMap::new(),
980 }
981}
982
983fn tid(value: &str) -> TypeId {
984 TypeId::new(value).expect("built-in type id is valid")
985}
986
987fn rid(value: &str) -> RepresentationId {
988 RepresentationId::new(value).expect("built-in representation id is valid")
989}
990
991#[cfg(test)]
992mod tests {
993 use std::collections::{BTreeMap, BTreeSet};
994
995 use super::*;
996 use crate::adapter::AdapterRegistry;
997 use crate::ids::SampleId;
998 use crate::model::DatasetSchema;
999 use crate::plan::DataPlanStepKind;
1000 use crate::planner::{plan_model_input, DataPlanRequest};
1001
1002 #[test]
1003 fn builtin_data_models_validate_and_have_unique_keys_and_representations() {
1004 let specs = builtin_data_model_specs();
1005 assert_eq!(specs.len(), BUILTIN_DATA_MODELS.len());
1006
1007 let mut keys = BTreeSet::new();
1008 let mut representations = BTreeSet::new();
1009 for spec in specs {
1010 assert!(keys.insert(spec.key.clone()), "duplicate key {}", spec.key);
1011 assert!(
1012 representations.insert(spec.representation.id.clone()),
1013 "duplicate representation {}",
1014 spec.representation.id
1015 );
1016 spec.representation.validate().unwrap();
1017 spec.source_descriptor(
1018 format!("source.{}", spec.key.replace('.', "_")),
1019 spec.key.clone(),
1020 "sample_id",
1021 SourceGranularity::PerSample,
1022 )
1023 .unwrap();
1024 }
1025 }
1026
1027 #[test]
1028 fn pasted_standardization_representations_are_present() {
1029 let representations = builtin_representations()
1030 .into_iter()
1031 .map(|representation| representation.id)
1032 .collect::<BTreeSet<_>>();
1033
1034 for expected in [
1035 REPRESENTATION_SIGNAL_1D,
1036 REPRESENTATION_SIGNAL_WITH_PROCESSINGS,
1037 REPRESENTATION_TABULAR_NUMERIC,
1038 REPRESENTATION_TABULAR_MIXED,
1039 REPRESENTATION_SERIES_MV,
1040 REPRESENTATION_VARIANT_MATRIX,
1041 REPRESENTATION_DOSAGE_MATRIX,
1042 REPRESENTATION_RGB_IMAGE,
1043 REPRESENTATION_GRAY_IMAGE,
1044 REPRESENTATION_CUBE_HWB,
1045 REPRESENTATION_FEATURE_BLOCK_SET,
1046 REPRESENTATION_SAMPLE_METADATA,
1047 REPRESENTATION_TARGET_NUMERIC,
1048 REPRESENTATION_TARGET_CATEGORICAL,
1049 REPRESENTATION_TARGET_NUMERIC_MATRIX,
1050 REPRESENTATION_TARGET_CATEGORICAL_MATRIX,
1051 REPRESENTATION_SEGMENTATION_MASK,
1052 REPRESENTATION_ROI_MASK,
1053 REPRESENTATION_MASS_SPECTRUM,
1054 REPRESENTATION_RAMAN_SIGNAL,
1055 REPRESENTATION_FTIR_SIGNAL,
1056 ] {
1057 assert!(
1058 representations.contains(&rid(expected)),
1059 "missing standardized representation {expected}"
1060 );
1061 }
1062 }
1063
1064 #[test]
1065 fn pasted_standardization_shapes_match_expected_axes() {
1066 let cube = hyperspectral_cube();
1067 assert_eq!(cube.id, rid(REPRESENTATION_CUBE_HWB));
1068 assert_eq!(axis_names(&cube), vec!["sample", "height", "width", "band"]);
1069
1070 let feature_blocks = feature_block_set();
1071 assert_eq!(feature_blocks.type_id, tid(TYPE_MULTI_BLOCK));
1072 assert_eq!(
1073 axis_names(&feature_blocks),
1074 vec!["sample", "block", "feature"]
1075 );
1076 assert!(feature_blocks.ragged);
1077
1078 let metadata = sample_metadata();
1079 assert_eq!(metadata.type_id, tid(TYPE_METADATA));
1080 assert_eq!(axis_names(&metadata), vec!["sample", "field"]);
1081
1082 for target in [target_numeric_matrix(), target_categorical_matrix()] {
1083 assert_eq!(target.type_id, tid(TYPE_TARGET));
1084 assert_eq!(axis_names(&target), vec!["sample", "target"]);
1085 assert_eq!(target.axes[1].kind, AxisKind::Target);
1086 }
1087
1088 for signal in [raman_signal(), ftir_signal()] {
1089 assert_eq!(signal.type_id, tid(TYPE_DENSE_SIGNAL));
1090 assert_eq!(axis_names(&signal), vec!["sample", "wavenumber"]);
1091 assert_eq!(signal.axes[1].kind, AxisKind::Wavenumber);
1092 }
1093
1094 let mass_spec = mass_spectrum();
1095 assert_eq!(mass_spec.type_id, tid(TYPE_MASS_SPEC));
1096 assert_eq!(axis_names(&mass_spec), vec!["sample", "mz"]);
1097 assert!(mass_spec.ragged);
1098 }
1099
1100 #[test]
1101 fn builtin_adapter_registry_validates() {
1102 let spec = builtin_adapter_registry_spec();
1103 AdapterRegistry::from_spec(spec).unwrap();
1104 }
1105
1106 #[test]
1107 fn common_builtin_models_can_plan_to_tabular_numeric() {
1108 let registry = AdapterRegistry::from_spec(builtin_adapter_registry_spec()).unwrap();
1109 let policy = default_builtin_planning_policy();
1110 let target_type = tid(TYPE_TABLE);
1111 let target_representation = rid(REPRESENTATION_TABULAR_NUMERIC);
1112
1113 for model in [
1114 BuiltinDataModel::NirsSignal1d,
1115 BuiltinDataModel::NirsSignalWithProcessings,
1116 BuiltinDataModel::RamanSignal,
1117 BuiltinDataModel::FtirSignal,
1118 BuiltinDataModel::TabularNumeric,
1119 BuiltinDataModel::TabularMixed,
1120 BuiltinDataModel::FeatureBlockSet,
1121 BuiltinDataModel::SeriesMultivariate,
1122 BuiltinDataModel::ClimateSeriesMultivariate,
1123 BuiltinDataModel::GenotypeVariantMatrix,
1124 BuiltinDataModel::GenotypeDosageMatrix,
1125 BuiltinDataModel::RgbImage,
1126 BuiltinDataModel::GrayImage,
1127 BuiltinDataModel::MultichannelImage,
1128 BuiltinDataModel::MultispectralImage,
1129 BuiltinDataModel::HyperspectralCube,
1130 BuiltinDataModel::SegmentationMask,
1131 BuiltinDataModel::RoiMask,
1132 BuiltinDataModel::SampleMetadata,
1133 BuiltinDataModel::MassSpectrum,
1134 BuiltinDataModel::TextRaw,
1135 BuiltinDataModel::TextTokenIds,
1136 ] {
1137 let representation = model.representation();
1138 let path = registry.find_path(
1139 &representation.type_id,
1140 &representation.id,
1141 &target_type,
1142 &target_representation,
1143 &policy,
1144 );
1145 assert!(
1146 path.path.is_some(),
1147 "no tabular path for {} ({}/{})",
1148 model.key(),
1149 representation.type_id,
1150 representation.id
1151 );
1152 }
1153 }
1154
1155 #[test]
1156 fn tabular_numeric_model_input_accepts_builtin_target() {
1157 let spec = tabular_numeric_model_input_spec();
1158 spec.validate().unwrap();
1159 let port = &spec.ports[0];
1160 assert_eq!(
1161 port.accepted_representations,
1162 vec![rid(REPRESENTATION_TABULAR_NUMERIC)]
1163 );
1164 assert_eq!(port.accepted_types, vec![tid(TYPE_TABLE)]);
1165 }
1166
1167 #[test]
1168 fn common_builtin_models_plan_model_input_to_tabular_numeric() {
1169 let registry = AdapterRegistry::from_spec(builtin_adapter_registry_spec()).unwrap();
1170 let model_input = tabular_numeric_model_input_spec();
1171
1172 for model in tabular_input_models() {
1173 let spec = model.spec();
1174 let source_id = format!("source.{}", spec.key.replace('.', "_"));
1175 let source = spec
1176 .source_descriptor(
1177 source_id,
1178 spec.key.clone(),
1179 "sample_id",
1180 SourceGranularity::PerSample,
1181 )
1182 .unwrap();
1183 let schema = DatasetSchema {
1184 dataset_id: format!("dataset.{}", spec.key.replace('.', "_")),
1185 sample_ids: vec![SampleId::new("S001").unwrap()],
1186 sources: vec![source],
1187 targets: BTreeMap::new(),
1188 metadata: BTreeMap::new(),
1189 metadata_schema: None,
1190 groups: Vec::new(),
1191 folds: Vec::new(),
1192 };
1193 let request = DataPlanRequest {
1194 id: format!("plan.{}", spec.key.replace('.', "_")),
1195 source_ids: None,
1196 planning_policy: default_builtin_planning_policy(),
1197 };
1198
1199 let plan = plan_model_input(&schema, &model_input, ®istry, &request).unwrap();
1200 plan.validate().unwrap();
1201 assert_eq!(
1202 plan.output_representation,
1203 rid(REPRESENTATION_TABULAR_NUMERIC)
1204 );
1205 assert!(
1206 plan.steps
1207 .iter()
1208 .any(|step| step.kind == DataPlanStepKind::Materialize),
1209 "{} plan did not materialize a source",
1210 spec.key
1211 );
1212 assert!(
1213 plan.steps
1214 .iter()
1215 .any(|step| step.kind == DataPlanStepKind::Join),
1216 "{} plan did not join into the model input port",
1217 spec.key
1218 );
1219 let has_adapt = plan
1220 .steps
1221 .iter()
1222 .any(|step| step.kind == DataPlanStepKind::Adapt);
1223 assert_eq!(
1224 has_adapt,
1225 model != BuiltinDataModel::TabularNumeric,
1226 "{} adapt-step expectation mismatch",
1227 spec.key
1228 );
1229 }
1230 }
1231
1232 fn axis_names(representation: &RepresentationSpec) -> Vec<&str> {
1233 representation
1234 .axes
1235 .iter()
1236 .map(|axis| axis.name.as_str())
1237 .collect()
1238 }
1239
1240 fn tabular_input_models() -> [BuiltinDataModel; 22] {
1241 [
1242 BuiltinDataModel::NirsSignal1d,
1243 BuiltinDataModel::NirsSignalWithProcessings,
1244 BuiltinDataModel::RamanSignal,
1245 BuiltinDataModel::FtirSignal,
1246 BuiltinDataModel::TabularNumeric,
1247 BuiltinDataModel::TabularMixed,
1248 BuiltinDataModel::FeatureBlockSet,
1249 BuiltinDataModel::SeriesMultivariate,
1250 BuiltinDataModel::ClimateSeriesMultivariate,
1251 BuiltinDataModel::GenotypeVariantMatrix,
1252 BuiltinDataModel::GenotypeDosageMatrix,
1253 BuiltinDataModel::RgbImage,
1254 BuiltinDataModel::GrayImage,
1255 BuiltinDataModel::MultichannelImage,
1256 BuiltinDataModel::MultispectralImage,
1257 BuiltinDataModel::HyperspectralCube,
1258 BuiltinDataModel::SegmentationMask,
1259 BuiltinDataModel::RoiMask,
1260 BuiltinDataModel::SampleMetadata,
1261 BuiltinDataModel::MassSpectrum,
1262 BuiltinDataModel::TextRaw,
1263 BuiltinDataModel::TextTokenIds,
1264 ]
1265 }
1266}