psyche_graphics/
obj.rs

1use obj_exporter::*;
2use psyche_core::brain::BrainActivityMap;
3use psyche_core::error::*;
4use psyche_core::neuron::Position;
5use psyche_core::Scalar;
6use std::io::Cursor;
7use std::iter::repeat;
8use std::string::FromUtf8Error;
9
10#[derive(Debug, Clone)]
11pub struct Config {
12    pub neurons: Option<Color>,
13    pub synapses: Option<Color>,
14    pub impulses: Option<Color>,
15    pub sensors: Option<Color>,
16    pub effectors: Option<Color>,
17    pub color_storage: ColorStorage,
18}
19
20impl Default for Config {
21    fn default() -> Self {
22        Self {
23            neurons: Some([255, 0, 255].into()),
24            synapses: Some([0, 0, 255].into()),
25            impulses: Some([192, 192, 255].into()),
26            sensors: Some([255, 255, 0].into()),
27            effectors: Some([128, 0, 0].into()),
28            color_storage: ColorStorage::Nowhere,
29        }
30    }
31}
32
33#[derive(Debug, Copy, Clone, PartialEq, Eq)]
34pub enum ColorStorage {
35    Nowhere,
36    Normals,
37    TexVertices,
38}
39
40impl Default for ColorStorage {
41    fn default() -> Self {
42        ColorStorage::Nowhere
43    }
44}
45
46/// (R, G, B)
47#[derive(Debug, Copy, Clone)]
48pub struct Color(u8, u8, u8);
49
50impl From<(u8, u8, u8)> for Color {
51    fn from(value: (u8, u8, u8)) -> Self {
52        Self(value.0, value.1, value.2)
53    }
54}
55
56impl From<[u8; 3]> for Color {
57    fn from(value: [u8; 3]) -> Self {
58        Self(value[0], value[1], value[2])
59    }
60}
61
62/// generates OBJ string from activity map.
63/// NOTE: Colors are stored either in vertices normals or texture vertices.
64pub fn generate_string(activity_map: &BrainActivityMap, config: &Config) -> Result<String> {
65    match String::from_utf8(generate(activity_map, config)?) {
66        Ok(s) => Ok(s),
67        Err(e) => Err(utf8_into_error(e)),
68    }
69}
70
71/// generates OBJ bytes from activity map.
72/// NOTE: Colors are stored either in vertices normals or texture vertices.
73pub fn generate(activity_map: &BrainActivityMap, config: &Config) -> Result<Vec<u8>> {
74    let mut objects = vec![];
75
76    if let Some(ref neurons) = config.neurons {
77        if !activity_map.neurons.is_empty() {
78            objects.push(Object {
79                name: "neurons".to_owned(),
80                vertices: activity_map
81                    .neurons
82                    .iter()
83                    .map(|p| Vertex {
84                        x: p.x,
85                        y: p.y,
86                        z: p.z,
87                    })
88                    .collect(),
89                tex_vertices: if config.color_storage == ColorStorage::TexVertices {
90                    let Color(r, g, b) = neurons;
91                    repeat(TVertex {
92                        u: Scalar::from(*r) / 255.0,
93                        v: Scalar::from(*g) / 255.0,
94                        w: Scalar::from(*b) / 255.0,
95                    })
96                    .take(activity_map.neurons.len())
97                    .collect()
98                } else {
99                    vec![]
100                },
101                normals: if config.color_storage == ColorStorage::Normals {
102                    let Color(r, g, b) = neurons;
103                    repeat(Vertex {
104                        x: Scalar::from(*r) / 255.0,
105                        y: Scalar::from(*g) / 255.0,
106                        z: Scalar::from(*b) / 255.0,
107                    })
108                    .take(activity_map.neurons.len())
109                    .collect()
110                } else {
111                    vec![]
112                },
113                geometry: vec![Geometry {
114                    material_name: None,
115                    shapes: activity_map
116                        .neurons
117                        .iter()
118                        .enumerate()
119                        .map(|(i, _)| Shape {
120                            primitive: Primitive::Point((i, None, None)),
121                            groups: vec![],
122                            smoothing_groups: vec![],
123                        })
124                        .collect(),
125                }],
126            });
127        }
128    }
129
130    if let Some(ref synapses) = config.synapses {
131        if !activity_map.connections.is_empty() {
132            objects.push(Object {
133                name: "synapses".to_owned(),
134                vertices: activity_map
135                    .connections
136                    .iter()
137                    .flat_map(|(f, t, _)| {
138                        vec![
139                            Vertex {
140                                x: f.x,
141                                y: f.y,
142                                z: f.z,
143                            },
144                            Vertex {
145                                x: t.x,
146                                y: t.y,
147                                z: t.z,
148                            },
149                        ]
150                    })
151                    .collect(),
152                tex_vertices: if config.color_storage == ColorStorage::TexVertices {
153                    let Color(r, g, b) = synapses;
154                    repeat(TVertex {
155                        u: Scalar::from(*r) / 255.0,
156                        v: Scalar::from(*g) / 255.0,
157                        w: Scalar::from(*b) / 255.0,
158                    })
159                    .take(activity_map.connections.len())
160                    .collect()
161                } else {
162                    vec![]
163                },
164                normals: if config.color_storage == ColorStorage::Normals {
165                    let Color(r, g, b) = synapses;
166                    repeat(Vertex {
167                        x: Scalar::from(*r) / 255.0,
168                        y: Scalar::from(*g) / 255.0,
169                        z: Scalar::from(*b) / 255.0,
170                    })
171                    .take(activity_map.connections.len())
172                    .collect()
173                } else {
174                    vec![]
175                },
176                geometry: vec![Geometry {
177                    material_name: None,
178                    shapes: activity_map
179                        .connections
180                        .iter()
181                        .enumerate()
182                        .map(|(i, _)| Shape {
183                            primitive: Primitive::Line(
184                                (i * 2, None, None),
185                                (i * 2 + 1, None, None),
186                            ),
187                            groups: vec![],
188                            smoothing_groups: vec![],
189                        })
190                        .collect(),
191                }],
192            });
193        }
194    }
195
196    if let Some(ref impulses) = config.impulses {
197        if !activity_map.impulses.is_empty() {
198            let positions = activity_map
199                .impulses
200                .iter()
201                .map(|(s, e, f)| lerp(*s, *e, *f))
202                .collect::<Vec<_>>();
203            objects.push(Object {
204                name: "impulses".to_owned(),
205                vertices: positions
206                    .iter()
207                    .map(|p| Vertex {
208                        x: p.x,
209                        y: p.y,
210                        z: p.z,
211                    })
212                    .collect(),
213                tex_vertices: if config.color_storage == ColorStorage::TexVertices {
214                    let Color(r, g, b) = impulses;
215                    repeat(TVertex {
216                        u: Scalar::from(*r) / 255.0,
217                        v: Scalar::from(*g) / 255.0,
218                        w: Scalar::from(*b) / 255.0,
219                    })
220                    .take(positions.len())
221                    .collect()
222                } else {
223                    vec![]
224                },
225                normals: if config.color_storage == ColorStorage::Normals {
226                    let Color(r, g, b) = impulses;
227                    repeat(Vertex {
228                        x: Scalar::from(*r) / 255.0,
229                        y: Scalar::from(*g) / 255.0,
230                        z: Scalar::from(*b) / 255.0,
231                    })
232                    .take(positions.len())
233                    .collect()
234                } else {
235                    vec![]
236                },
237                geometry: vec![Geometry {
238                    material_name: None,
239                    shapes: positions
240                        .iter()
241                        .enumerate()
242                        .map(|(i, _)| Shape {
243                            primitive: Primitive::Point((i, None, None)),
244                            groups: vec![],
245                            smoothing_groups: vec![],
246                        })
247                        .collect(),
248                }],
249            });
250        }
251    }
252
253    if let Some(ref sensors) = config.sensors {
254        if !activity_map.sensors.is_empty() {
255            objects.push(Object {
256                name: "sensors".to_owned(),
257                vertices: activity_map
258                    .sensors
259                    .iter()
260                    .map(|p| Vertex {
261                        x: p.x,
262                        y: p.y,
263                        z: p.z,
264                    })
265                    .collect(),
266                tex_vertices: if config.color_storage == ColorStorage::TexVertices {
267                    let Color(r, g, b) = sensors;
268                    repeat(TVertex {
269                        u: Scalar::from(*r) / 255.0,
270                        v: Scalar::from(*g) / 255.0,
271                        w: Scalar::from(*b) / 255.0,
272                    })
273                    .take(activity_map.sensors.len())
274                    .collect()
275                } else {
276                    vec![]
277                },
278                normals: if config.color_storage == ColorStorage::Normals {
279                    let Color(r, g, b) = sensors;
280                    repeat(Vertex {
281                        x: Scalar::from(*r) / 255.0,
282                        y: Scalar::from(*g) / 255.0,
283                        z: Scalar::from(*b) / 255.0,
284                    })
285                    .take(activity_map.sensors.len())
286                    .collect()
287                } else {
288                    vec![]
289                },
290                geometry: vec![Geometry {
291                    material_name: None,
292                    shapes: activity_map
293                        .sensors
294                        .iter()
295                        .enumerate()
296                        .map(|(i, _)| Shape {
297                            primitive: Primitive::Point((i, None, None)),
298                            groups: vec![],
299                            smoothing_groups: vec![],
300                        })
301                        .collect(),
302                }],
303            });
304        }
305    }
306
307    if let Some(ref effectors) = config.effectors {
308        if !activity_map.effectors.is_empty() {
309            objects.push(Object {
310                name: "effectors".to_owned(),
311                vertices: activity_map
312                    .effectors
313                    .iter()
314                    .map(|p| Vertex {
315                        x: p.x,
316                        y: p.y,
317                        z: p.z,
318                    })
319                    .collect(),
320                tex_vertices: if config.color_storage == ColorStorage::TexVertices {
321                    let Color(r, g, b) = effectors;
322                    repeat(TVertex {
323                        u: Scalar::from(*r) / 255.0,
324                        v: Scalar::from(*g) / 255.0,
325                        w: Scalar::from(*b) / 255.0,
326                    })
327                    .take(activity_map.effectors.len())
328                    .collect()
329                } else {
330                    vec![]
331                },
332                normals: if config.color_storage == ColorStorage::Normals {
333                    let Color(r, g, b) = effectors;
334                    repeat(Vertex {
335                        x: Scalar::from(*r) / 255.0,
336                        y: Scalar::from(*g) / 255.0,
337                        z: Scalar::from(*b) / 255.0,
338                    })
339                    .take(activity_map.effectors.len())
340                    .collect()
341                } else {
342                    vec![]
343                },
344                geometry: vec![Geometry {
345                    material_name: None,
346                    shapes: activity_map
347                        .effectors
348                        .iter()
349                        .enumerate()
350                        .map(|(i, _)| Shape {
351                            primitive: Primitive::Point((i, None, None)),
352                            groups: vec![],
353                            smoothing_groups: vec![],
354                        })
355                        .collect(),
356                }],
357            });
358        }
359    }
360
361    let set = ObjSet {
362        material_library: None,
363        objects,
364    };
365    let mut cursor = Cursor::new(vec![]);
366    export(&set, &mut cursor)?;
367    Ok(cursor.into_inner())
368}
369
370fn lerp(start: Position, end: Position, factor: Scalar) -> Position {
371    let factor = factor.max(0.0).min(1.0);
372    Position {
373        x: (end.x - start.x) * factor + start.x,
374        y: (end.y - start.y) * factor + start.y,
375        z: (end.z - start.z) * factor + start.z,
376    }
377}
378
379fn utf8_into_error(error: FromUtf8Error) -> Error {
380    Error::simple(format!("{}", error))
381}