initiative_core/world/
demographics.rs

1use super::npc::{Ethnicity, Species};
2use initiative_macros::From;
3
4use rand::distributions::WeightedIndex;
5use rand::prelude::*;
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::iter;
9
10type GroupMap = HashMap<(Species, Ethnicity), u64>;
11type GroupMapSerialized = Vec<(Species, Ethnicity, u64)>;
12
13#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
14pub struct Demographics {
15    groups: GroupMapWrapper,
16}
17
18#[derive(Clone, Debug, Deserialize, Eq, From, PartialEq, Serialize)]
19#[serde(from = "GroupMapSerialized", into = "GroupMapSerialized")]
20struct GroupMapWrapper(GroupMap);
21
22impl Demographics {
23    pub fn new(groups: GroupMap) -> Self {
24        Self {
25            groups: GroupMapWrapper(groups),
26        }
27    }
28
29    pub fn shift_species(&self, species: &Species, amount: f64) -> Self {
30        self.shift_by(
31            |s, _| s == species,
32            amount,
33            (*species, species.default_ethnicity()),
34        )
35    }
36
37    pub fn only_species(&self, species: &Species) -> Self {
38        self.shift_species(species, 1.)
39    }
40
41    pub fn shift_ethnicity(&self, ethnicity: &Ethnicity, amount: f64) -> Self {
42        self.shift_by(
43            |_, e| e == ethnicity,
44            amount,
45            (ethnicity.default_species(), *ethnicity),
46        )
47    }
48
49    pub fn only_ethnicity(&self, ethnicity: &Ethnicity) -> Self {
50        self.shift_ethnicity(ethnicity, 1.)
51    }
52
53    pub fn shift_species_ethnicity(
54        &self,
55        species: &Species,
56        ethnicity: &Ethnicity,
57        amount: f64,
58    ) -> Self {
59        self.shift_by(
60            |s, e| s == species && e == ethnicity,
61            amount,
62            (*species, *ethnicity),
63        )
64    }
65
66    pub fn only_species_ethnicity(&self, species: &Species, ethnicity: &Ethnicity) -> Self {
67        self.shift_species_ethnicity(species, ethnicity, 1.)
68    }
69
70    pub fn gen_species_ethnicity(&self, rng: &mut impl Rng) -> (Species, Ethnicity) {
71        if self.groups().is_empty() {
72            (Species::Human, Species::Human.default_ethnicity())
73        } else {
74            let (groups, weights): (Vec<(Species, Ethnicity)>, Vec<u64>) =
75                self.groups().iter().unzip();
76            let dist = WeightedIndex::new(weights).unwrap();
77            groups[dist.sample(rng)]
78        }
79    }
80
81    fn shift_by<F: Fn(&Species, &Ethnicity) -> bool>(
82        &self,
83        f: F,
84        amount: f64,
85        default: (Species, Ethnicity),
86    ) -> Self {
87        if !(0. ..=1.).contains(&amount) {
88            panic!("Invalid input: {}", amount);
89        }
90
91        let population: u64 = self.groups().values().sum();
92        let species_population: u64 = self
93            .groups()
94            .iter()
95            .filter_map(|((s, e), n)| if f(s, e) { Some(n) } else { None })
96            .sum();
97
98        let groups: GroupMap = if species_population > 0 {
99            self.groups()
100                .iter()
101                .map(|((s, e), &v)| {
102                    (
103                        (*s, *e),
104                        if f(s, e) {
105                            (v as f64 * (1. - amount)
106                                + (v as f64 * amount * population as f64
107                                    / species_population as f64))
108                                .round() as u64
109                        } else {
110                            (v as f64 * (1. - amount)).round() as u64
111                        },
112                    )
113                })
114                .filter(|(_, v)| *v > 0)
115                .collect()
116        } else {
117            self.groups()
118                .iter()
119                .map(|(&k, &v)| (k, (v as f64 * (1. - amount)).round() as u64))
120                .chain(iter::once((
121                    default,
122                    (population as f64 * amount).round() as u64,
123                )))
124                .filter(|(_, v)| *v > 0)
125                .collect()
126        };
127
128        Self::new(groups)
129    }
130
131    fn groups(&self) -> &GroupMap {
132        &self.groups.0
133    }
134}
135
136impl Default for Demographics {
137    fn default() -> Self {
138        let mut groups = HashMap::new();
139        groups.insert((Species::Human, Ethnicity::Human), 1_020_000);
140        groups.insert((Species::HalfElf, Ethnicity::Elvish), 320_000);
141        groups.insert((Species::Elf, Ethnicity::Elvish), 220_000);
142        groups.insert((Species::Gnome, Ethnicity::Gnomish), 220_000);
143        groups.insert((Species::Halfling, Ethnicity::Halfling), 100_000);
144        // groups.insert(Species::Shifter, 60_000);
145        // groups.insert(Species::Changeling, 40_000);
146
147        Self::new(groups)
148    }
149}
150
151impl From<GroupMapWrapper> for GroupMap {
152    fn from(value: GroupMapWrapper) -> Self {
153        value.0
154    }
155}
156
157impl From<GroupMapSerialized> for GroupMapWrapper {
158    fn from(value: GroupMapSerialized) -> Self {
159        Self(value.into_iter().map(|(a, b, c)| ((a, b), c)).collect())
160    }
161}
162
163impl From<GroupMapWrapper> for GroupMapSerialized {
164    fn from(value: GroupMapWrapper) -> Self {
165        value.0.into_iter().map(|((a, b), c)| (a, b, c)).collect()
166    }
167}
168
169#[cfg(test)]
170mod test {
171    use super::*;
172
173    #[test]
174    fn shift_species_test_existing() {
175        let demographics = demographics().shift_species(&Species::Human, 0.3);
176
177        assert_eq!(3, demographics.groups().len());
178        assert_eq!(
179            Some(&39),
180            demographics
181                .groups()
182                .get(&(Species::Human, Ethnicity::Human)),
183            "{:?}",
184            demographics
185        );
186        assert_eq!(
187            Some(&26),
188            demographics
189                .groups()
190                .get(&(Species::Human, Ethnicity::Gnomish))
191        );
192        assert_eq!(
193            Some(&35),
194            demographics
195                .groups()
196                .get(&(Species::Gnome, Ethnicity::Gnomish))
197        );
198    }
199
200    #[test]
201    fn shift_ethnicity_test_existing() {
202        let demographics = demographics().shift_ethnicity(&Ethnicity::Gnomish, 0.3);
203
204        assert_eq!(3, demographics.groups().len());
205        assert_eq!(
206            Some(&21),
207            demographics
208                .groups()
209                .get(&(Species::Human, Ethnicity::Human)),
210            "{:?}",
211            demographics
212        );
213        assert_eq!(
214            Some(&23),
215            demographics
216                .groups()
217                .get(&(Species::Human, Ethnicity::Gnomish))
218        );
219        assert_eq!(
220            Some(&56),
221            demographics
222                .groups()
223                .get(&(Species::Gnome, Ethnicity::Gnomish))
224        );
225    }
226
227    #[test]
228    fn shift_species_ethnicity_test_existing() {
229        let demographics =
230            demographics().shift_species_ethnicity(&Species::Gnome, &Ethnicity::Gnomish, 0.3);
231
232        assert_eq!(3, demographics.groups().len());
233        assert_eq!(
234            Some(&21),
235            demographics
236                .groups()
237                .get(&(Species::Human, Ethnicity::Human)),
238            "{:?}",
239            demographics
240        );
241        assert_eq!(
242            Some(&14),
243            demographics
244                .groups()
245                .get(&(Species::Human, Ethnicity::Gnomish))
246        );
247        assert_eq!(
248            Some(&65),
249            demographics
250                .groups()
251                .get(&(Species::Gnome, Ethnicity::Gnomish))
252        );
253    }
254
255    #[test]
256    fn shift_species_test_new() {
257        let mut groups = HashMap::with_capacity(1);
258        groups.insert((Species::Human, Ethnicity::Human), 100);
259        let demographics = Demographics::new(groups).shift_species(&Species::Gnome, 0.4);
260
261        assert_eq!(2, demographics.groups().len());
262        assert_eq!(
263            Some(&60),
264            demographics
265                .groups()
266                .get(&(Species::Human, Ethnicity::Human))
267        );
268        assert_eq!(
269            Some(&40),
270            demographics
271                .groups()
272                .get(&(Species::Gnome, Ethnicity::Gnomish))
273        );
274    }
275
276    #[test]
277    fn shift_ethnicity_test_new() {
278        let mut groups = HashMap::with_capacity(1);
279        groups.insert((Species::Human, Ethnicity::Human), 100);
280        let demographics = Demographics::new(groups).shift_ethnicity(&Ethnicity::Gnomish, 0.4);
281
282        assert_eq!(2, demographics.groups().len());
283        assert_eq!(
284            Some(&60),
285            demographics
286                .groups()
287                .get(&(Species::Human, Ethnicity::Human))
288        );
289        assert_eq!(
290            Some(&40),
291            demographics
292                .groups()
293                .get(&(Species::Gnome, Ethnicity::Gnomish))
294        );
295    }
296
297    #[test]
298    fn shift_species_ethnicity_test_new() {
299        let mut groups = HashMap::with_capacity(1);
300        groups.insert((Species::Human, Ethnicity::Human), 100);
301        let demographics = Demographics::new(groups).shift_species_ethnicity(
302            &Species::Gnome,
303            &Ethnicity::Gnomish,
304            0.4,
305        );
306
307        assert_eq!(2, demographics.groups().len());
308        assert_eq!(
309            Some(&60),
310            demographics
311                .groups()
312                .get(&(Species::Human, Ethnicity::Human))
313        );
314        assert_eq!(
315            Some(&40),
316            demographics
317                .groups()
318                .get(&(Species::Gnome, Ethnicity::Gnomish))
319        );
320    }
321
322    #[test]
323    fn only_species_test() {
324        let demographics = demographics().only_species(&Species::Human);
325
326        assert_eq!(2, demographics.groups().len());
327        assert_eq!(
328            Some(&60),
329            demographics
330                .groups()
331                .get(&(Species::Human, Ethnicity::Human))
332        );
333        assert_eq!(
334            Some(&40),
335            demographics
336                .groups()
337                .get(&(Species::Human, Ethnicity::Gnomish))
338        );
339    }
340
341    #[test]
342    fn only_ethnicity_test() {
343        let demographics = demographics().only_ethnicity(&Ethnicity::Gnomish);
344
345        assert_eq!(2, demographics.groups().len());
346        assert_eq!(
347            Some(&29),
348            demographics
349                .groups()
350                .get(&(Species::Human, Ethnicity::Gnomish))
351        );
352        assert_eq!(
353            Some(&71),
354            demographics
355                .groups()
356                .get(&(Species::Gnome, Ethnicity::Gnomish))
357        );
358    }
359
360    #[test]
361    fn only_species_ethnicity_test() {
362        let demographics =
363            demographics().only_species_ethnicity(&Species::Gnome, &Ethnicity::Gnomish);
364
365        assert_eq!(1, demographics.groups().len());
366        assert_eq!(
367            Some(&100),
368            demographics
369                .groups()
370                .get(&(Species::Gnome, Ethnicity::Gnomish))
371        );
372    }
373
374    #[test]
375    fn gen_species_ethnicity_test() {
376        let mut groups = HashMap::new();
377        groups.insert((Species::Human, Ethnicity::Human), 50);
378        groups.insert((Species::Gnome, Ethnicity::Gnomish), 50);
379        let demographics = Demographics::new(groups);
380
381        let mut rng = SmallRng::seed_from_u64(0);
382        let mut counts: HashMap<(Species, Ethnicity), u8> = HashMap::with_capacity(2);
383
384        for i in 0..10 {
385            let species_ethnicity = &demographics.gen_species_ethnicity(&mut rng);
386            *counts.entry(*species_ethnicity).or_default() += 1;
387            println!("{}: {:?}", i, counts);
388        }
389
390        assert_eq!(Some(&5), counts.get(&(Species::Human, Ethnicity::Human)));
391        assert_eq!(Some(&5), counts.get(&(Species::Gnome, Ethnicity::Gnomish)));
392    }
393
394    #[test]
395    fn demographics_serialize_deserialize_test() {
396        let demographics = demographics();
397
398        let value: Demographics = serde_json::from_str(
399            r#"{"groups":[["gnome","gnomish",50],["human","gnomish",20],["human","human",30]]}"#,
400        )
401        .unwrap();
402        assert_eq!(demographics, value);
403
404        // Ordering is not guaranteed in hash maps, so we have to fudge it with a round-trip.
405        assert_eq!(
406            demographics,
407            serde_json::from_str(&serde_json::to_string(&demographics).unwrap()).unwrap(),
408        );
409    }
410
411    fn demographics() -> Demographics {
412        let mut groups = HashMap::with_capacity(3);
413        groups.insert((Species::Human, Ethnicity::Human), 30);
414        groups.insert((Species::Human, Ethnicity::Gnomish), 20);
415        groups.insert((Species::Gnome, Ethnicity::Gnomish), 50);
416        Demographics::new(groups)
417    }
418}
initiative_core/world/demographics.rs

initiative_core/world/
demographics.rs
