Skip to content

Commit 7e09c81

Browse files
committed
Refactor MeshAttributes (support for non-owned data)
1 parent ba2160b commit 7e09c81

File tree

10 files changed

+193
-145
lines changed

10 files changed

+193
-145
lines changed

pysplashsurf/src/mesh.rs

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ use pyo3_stub_gen::derive::*;
1010
use splashsurf_lib::{
1111
Real,
1212
mesh::{
13-
AttributeData, Mesh3d, MeshAttribute, MeshWithData, MixedTriQuadMesh3d, TriMesh3d,
14-
TriangleOrQuadCell,
13+
Mesh3d, MeshWithData, MixedTriQuadMesh3d, OwnedAttributeData, OwnedMeshAttribute,
14+
TriMesh3d, TriangleOrQuadCell,
1515
},
1616
nalgebra::{Unit, Vector3},
1717
};
@@ -20,7 +20,7 @@ use crate::aabb::{Aabb3dF32, Aabb3dF64};
2020

2121
fn get_attribute_with_name<'py, R: Real + Element>(
2222
py: Python<'py>,
23-
attrs: &[MeshAttribute<R>],
23+
attrs: &[OwnedMeshAttribute<R>],
2424
name: &str,
2525
) -> PyResult<PyObject>
2626
where
@@ -29,10 +29,14 @@ where
2929
let elem = attrs.iter().filter(|x| x.name == name).next();
3030
match elem {
3131
Some(attr) => match attr.data.clone() {
32-
AttributeData::ScalarU64(res) => Ok(res.into_pyobject(py).unwrap().into()),
33-
AttributeData::ScalarReal(res) => Ok(res.into_pyobject(py).unwrap().into()),
34-
AttributeData::Vector3Real(res) => {
35-
let flattened: Vec<R> = bytemuck::cast_vec(res);
32+
OwnedAttributeData::ScalarU64(res) => {
33+
Ok(res.into_owned().into_pyobject(py).unwrap().into())
34+
}
35+
OwnedAttributeData::ScalarReal(res) => {
36+
Ok(res.into_owned().into_pyobject(py).unwrap().into())
37+
}
38+
OwnedAttributeData::Vector3Real(res) => {
39+
let flattened: Vec<R> = bytemuck::cast_vec(res.into_owned());
3640
let res: Array2<R> =
3741
Array2::from_shape_vec((flattened.len() / 3, 3), flattened).unwrap();
3842
Ok(res.into_pyarray(py).into_bound_py_any(py).unwrap().into())
@@ -46,8 +50,8 @@ where
4650
}
4751

4852
fn add_attribute_with_name<'py, R: Real + Element>(
49-
attrs: &mut Vec<MeshAttribute<R>>,
50-
attribute: MeshAttribute<R>,
53+
attrs: &mut Vec<OwnedMeshAttribute<R>>,
54+
attribute: OwnedMeshAttribute<R>,
5155
) -> PyResult<()> {
5256
let elem = attrs.iter().filter(|x| x.name == attribute.name).next();
5357
match elem {
@@ -119,7 +123,7 @@ macro_rules! create_mesh_data_interface {
119123
) -> PyResult<()> {
120124
add_attribute_with_name::<$type>(
121125
&mut self.inner.point_attributes,
122-
MeshAttribute::new(name, AttributeData::ScalarU64(data)),
126+
OwnedMeshAttribute::new(name, OwnedAttributeData::ScalarU64(data.into())),
123127
)
124128
}
125129

@@ -130,7 +134,7 @@ macro_rules! create_mesh_data_interface {
130134
) -> PyResult<()> {
131135
add_attribute_with_name::<$type>(
132136
&mut self.inner.point_attributes,
133-
MeshAttribute::new(name, AttributeData::ScalarReal(data)),
137+
OwnedMeshAttribute::new(name, OwnedAttributeData::ScalarReal(data.into())),
134138
)
135139
}
136140

@@ -145,7 +149,10 @@ macro_rules! create_mesh_data_interface {
145149

146150
add_attribute_with_name::<$type>(
147151
&mut self.inner.point_attributes,
148-
MeshAttribute::new(name, AttributeData::Vector3Real(data.to_vec())),
152+
OwnedMeshAttribute::new(
153+
name,
154+
OwnedAttributeData::Vector3Real(data.to_vec().into()),
155+
),
149156
)
150157
}
151158

@@ -156,7 +163,7 @@ macro_rules! create_mesh_data_interface {
156163
) -> PyResult<()> {
157164
add_attribute_with_name::<$type>(
158165
&mut self.inner.cell_attributes,
159-
MeshAttribute::new(name, AttributeData::ScalarU64(data)),
166+
OwnedMeshAttribute::new(name, OwnedAttributeData::ScalarU64(data.into())),
160167
)
161168
}
162169

@@ -167,7 +174,7 @@ macro_rules! create_mesh_data_interface {
167174
) -> PyResult<()> {
168175
add_attribute_with_name::<$type>(
169176
&mut self.inner.cell_attributes,
170-
MeshAttribute::new(name, AttributeData::ScalarReal(data)),
177+
OwnedMeshAttribute::new(name, OwnedAttributeData::ScalarReal(data.into())),
171178
)
172179
}
173180

@@ -182,7 +189,10 @@ macro_rules! create_mesh_data_interface {
182189

183190
add_attribute_with_name::<$type>(
184191
&mut self.inner.cell_attributes,
185-
MeshAttribute::new(name, AttributeData::Vector3Real(data.to_vec())),
192+
OwnedMeshAttribute::new(
193+
name,
194+
OwnedAttributeData::Vector3Real(data.to_vec().into()),
195+
),
186196
)
187197
}
188198

pysplashsurf/src/pipeline.rs

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,13 @@ use pyo3::{
1212
};
1313
use splashsurf_lib::{
1414
Aabb3d, GridDecompositionParameters, Index, Real, SpatialDecomposition,
15-
mesh::{AttributeData, MeshAttribute},
15+
mesh::{OwnedAttributeData, OwnedMeshAttribute},
1616
nalgebra::Vector3,
1717
};
1818

1919
fn reconstruction_pipeline_generic<I: Index, R: Real>(
2020
particle_positions: &[Vector3<R>],
21-
attributes: Vec<MeshAttribute<R>>,
21+
attributes: Vec<OwnedMeshAttribute<R>>,
2222
particle_radius: R,
2323
rest_density: R,
2424
smoothing_length: R,
@@ -134,8 +134,8 @@ fn reconstruction_pipeline_generic<I: Index, R: Real>(
134134

135135
fn attrs_conversion<R: Real + Element>(
136136
attributes_to_interpolate: Bound<PyDict>,
137-
) -> Vec<MeshAttribute<R>> {
138-
let mut attrs: Vec<MeshAttribute<R>> = Vec::new();
137+
) -> Vec<OwnedMeshAttribute<R>> {
138+
let mut attrs: Vec<OwnedMeshAttribute<R>> = Vec::new();
139139
for (key, value) in attributes_to_interpolate.iter() {
140140
let key_str: String = key
141141
.downcast::<PyString>()
@@ -150,7 +150,8 @@ fn attrs_conversion<R: Real + Element>(
150150
.as_slice()
151151
.unwrap()
152152
.to_vec();
153-
let mesh_attr = MeshAttribute::new(key_str, AttributeData::ScalarU64(value));
153+
let mesh_attr =
154+
OwnedMeshAttribute::new(key_str, OwnedAttributeData::ScalarU64(value.into()));
154155
attrs.push(mesh_attr);
155156
} else if let Ok(value) = value.downcast::<PyArray1<R>>() {
156157
let value: Vec<R> = value
@@ -159,16 +160,19 @@ fn attrs_conversion<R: Real + Element>(
159160
.as_slice()
160161
.unwrap()
161162
.to_vec();
162-
let mesh_attr = MeshAttribute::new(key_str, AttributeData::ScalarReal(value));
163+
let mesh_attr =
164+
OwnedMeshAttribute::new(key_str, OwnedAttributeData::ScalarReal(value.into()));
163165
attrs.push(mesh_attr);
164166
} else if let Ok(value) = value.downcast::<PyArray2<R>>() {
165167
let value: PyReadonlyArray2<R> = value.extract().unwrap();
166168

167169
let value_slice = value.as_slice().unwrap();
168170
let value_slice: &[Vector3<R>] = bytemuck::cast_slice(value_slice);
169171

170-
let mesh_attr =
171-
MeshAttribute::new(key_str, AttributeData::Vector3Real(value_slice.to_vec()));
172+
let mesh_attr = OwnedMeshAttribute::new(
173+
key_str,
174+
OwnedAttributeData::Vector3Real(value_slice.to_vec().into()),
175+
);
172176
attrs.push(mesh_attr);
173177
} else {
174178
println!("Couldnt downcast attribute {} to valid type", &key_str);

splashsurf/src/io.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use anyhow::{Context, anyhow};
33
use log::{info, warn};
44
use splashsurf_lib::Real;
55
use splashsurf_lib::mesh::{
6-
IntoVtkUnstructuredGridPiece, Mesh3d, MeshAttribute, MeshWithData, TriMesh3d,
6+
IntoVtkUnstructuredGridPiece, Mesh3d, MeshWithData, OwnedMeshAttribute, TriMesh3d,
77
};
88
use splashsurf_lib::nalgebra::Vector3;
99
use splashsurf_lib::{io, profile};
@@ -69,7 +69,7 @@ pub fn read_particle_positions_with_attributes<R: Real, P: AsRef<Path>>(
6969
input_file: P,
7070
attribute_names: &[String],
7171
format_params: &InputFormatParameters,
72-
) -> Result<(Vec<Vector3<R>>, Vec<MeshAttribute<R>>), anyhow::Error> {
72+
) -> Result<(Vec<Vector3<R>>, Vec<OwnedMeshAttribute<R>>), anyhow::Error> {
7373
if attribute_names.is_empty() {
7474
return read_particle_positions(input_file, format_params).map(|p| (p, Vec::new()));
7575
}

splashsurf/src/reconstruct.rs

Lines changed: 48 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use indicatif::{ProgressBar, ProgressStyle};
99
use log::{error, info, warn};
1010
use rayon::prelude::*;
1111
use splashsurf_lib::mesh::{
12-
AttributeData, Mesh3d, MeshAttribute, MeshWithData, MixedTriQuadMesh3d, TriMesh3d,
12+
Mesh3d, MeshWithData, MixedTriQuadMesh3d, OwnedAttributeData, OwnedMeshAttribute, TriMesh3d,
1313
};
1414
use splashsurf_lib::nalgebra::{Unit, Vector3};
1515
use splashsurf_lib::sph_interpolation::SphInterpolator;
@@ -1009,7 +1009,7 @@ pub(crate) fn reconstruction_pipeline_from_args(
10091009
/// in the post-processing parameters (see [`ReconstructionPostprocessingParameters`]).
10101010
pub fn reconstruction_pipeline<I: Index, R: Real>(
10111011
particle_positions: &[Vector3<R>],
1012-
attributes: &[MeshAttribute<R>],
1012+
attributes: &[OwnedMeshAttribute<R>],
10131013
params: &splashsurf_lib::Parameters<R>,
10141014
postprocessing: &ReconstructionPostprocessingParameters,
10151015
) -> Result<ReconstructionResult<I, R>, anyhow::Error> {
@@ -1220,15 +1220,19 @@ pub fn reconstruction_pipeline<I: Index, R: Real>(
12201220

12211221
if postprocessing.output_mesh_smoothing_weights {
12221222
// Raw distance-weighted number of neighbors value per vertex (can be used to determine normalization value)
1223-
mesh_with_data.point_attributes.push(MeshAttribute::new(
1224-
"wnn".to_string(),
1225-
AttributeData::ScalarReal(vertex_weighted_num_neighbors),
1226-
));
1223+
mesh_with_data
1224+
.point_attributes
1225+
.push(OwnedMeshAttribute::new(
1226+
"wnn".to_string(),
1227+
OwnedAttributeData::ScalarReal(vertex_weighted_num_neighbors.into()),
1228+
));
12271229
// Final smoothing weights per vertex
1228-
mesh_with_data.point_attributes.push(MeshAttribute::new(
1229-
"sw".to_string(),
1230-
AttributeData::ScalarReal(smoothing_weights.clone()),
1231-
));
1230+
mesh_with_data
1231+
.point_attributes
1232+
.push(OwnedMeshAttribute::new(
1233+
"sw".to_string(),
1234+
OwnedAttributeData::ScalarReal(smoothing_weights.clone().into()),
1235+
));
12321236
}
12331237

12341238
smoothing_weights
@@ -1296,21 +1300,27 @@ pub fn reconstruction_pipeline<I: Index, R: Real>(
12961300
smoothing_iters,
12971301
);
12981302

1299-
mesh_with_data.point_attributes.push(MeshAttribute::new(
1300-
"normals".to_string(),
1301-
AttributeData::Vector3Real(smoothed_normals),
1302-
));
1303-
if postprocessing.output_raw_normals {
1304-
mesh_with_data.point_attributes.push(MeshAttribute::new(
1305-
"raw_normals".to_string(),
1306-
AttributeData::Vector3Real(normals),
1303+
mesh_with_data
1304+
.point_attributes
1305+
.push(OwnedMeshAttribute::new(
1306+
"normals".to_string(),
1307+
OwnedAttributeData::Vector3Real(smoothed_normals.into()),
13071308
));
1309+
if postprocessing.output_raw_normals {
1310+
mesh_with_data
1311+
.point_attributes
1312+
.push(OwnedMeshAttribute::new(
1313+
"raw_normals".to_string(),
1314+
OwnedAttributeData::Vector3Real(normals.into()),
1315+
));
13081316
}
13091317
} else {
1310-
mesh_with_data.point_attributes.push(MeshAttribute::new(
1311-
"normals".to_string(),
1312-
AttributeData::Vector3Real(normals),
1313-
));
1318+
mesh_with_data
1319+
.point_attributes
1320+
.push(OwnedMeshAttribute::new(
1321+
"normals".to_string(),
1322+
OwnedAttributeData::Vector3Real(normals.into()),
1323+
));
13141324
}
13151325
}
13161326

@@ -1330,29 +1340,33 @@ pub fn reconstruction_pipeline<I: Index, R: Real>(
13301340

13311341
let particles_inside = reconstruction.particle_inside_aabb().map(Vec::as_slice);
13321342
match &attribute.data {
1333-
AttributeData::ScalarReal(values) => {
1334-
let filtered_values = filtered_quantity(values, particles_inside);
1343+
OwnedAttributeData::ScalarReal(values) => {
1344+
let filtered_values = filtered_quantity(&values, particles_inside);
13351345
let interpolated_values = interpolator.interpolate_scalar_quantity(
13361346
&filtered_values,
13371347
mesh_with_data.vertices(),
13381348
true,
13391349
);
1340-
mesh_with_data.point_attributes.push(MeshAttribute::new(
1341-
attribute.name.clone(),
1342-
AttributeData::ScalarReal(interpolated_values),
1343-
));
1350+
mesh_with_data
1351+
.point_attributes
1352+
.push(OwnedMeshAttribute::new(
1353+
attribute.name.clone(),
1354+
OwnedAttributeData::ScalarReal(interpolated_values.into()),
1355+
));
13441356
}
1345-
AttributeData::Vector3Real(values) => {
1346-
let filtered_values = filtered_quantity(values, particles_inside);
1357+
OwnedAttributeData::Vector3Real(values) => {
1358+
let filtered_values = filtered_quantity(&values, particles_inside);
13471359
let interpolated_values = interpolator.interpolate_vector_quantity(
13481360
&filtered_values,
13491361
mesh_with_data.vertices(),
13501362
true,
13511363
);
1352-
mesh_with_data.point_attributes.push(MeshAttribute::new(
1353-
attribute.name.clone(),
1354-
AttributeData::Vector3Real(interpolated_values),
1355-
));
1364+
mesh_with_data
1365+
.point_attributes
1366+
.push(OwnedMeshAttribute::new(
1367+
attribute.name.clone(),
1368+
OwnedAttributeData::Vector3Real(interpolated_values.into()),
1369+
));
13561370
}
13571371
_ => unimplemented!("Interpolation of this attribute type not implemented"),
13581372
}

splashsurf_lib/src/density_map.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
2121
use crate::aabb::Aabb3d;
2222
use crate::kernel::DiscreteSquaredDistanceCubicKernel;
23-
use crate::mesh::{HexMesh3d, MeshAttribute, MeshWithData};
23+
use crate::mesh::{HexMesh3d, MeshWithData, OwnedMeshAttribute};
2424
use crate::neighborhood_search::NeighborhoodList;
2525
use crate::uniform_grid::UniformGrid;
2626
use crate::utils::{ChunkSize, ParallelPolicy};
@@ -780,7 +780,7 @@ pub fn sparse_density_map_to_hex_mesh<I: Index, R: Real>(
780780
]);
781781
}
782782

783-
MeshWithData::new(mesh).with_point_data(MeshAttribute::new_real_scalar(
783+
MeshWithData::new(mesh).with_point_data(OwnedMeshAttribute::new_real_scalar(
784784
"density".to_string(),
785785
values,
786786
))

0 commit comments

Comments
 (0)