Skip to content
Open
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 125 additions & 20 deletions src/atlas4py/_atlas4py.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -231,19 +231,19 @@ NB_MODULE( _atlas4py, m ) {
.def_prop_ro( "lon", nb::overload_cast<>( &PointLonLat::lon, nb::const_ ) )
.def_prop_ro( "lat", nb::overload_cast<>( &PointLonLat::lat, nb::const_ ) )
.def( "__repr__", []( PointLonLat const& p ) {
return "_atlas4py.PointLonLat(lon=" + std::to_string( p.lon() ) + ", lat=" + std::to_string( p.lat() ) + ")";
return "atlas4py.PointLonLat(lon="_s + nb::str( nb::float_( p.lon() ) ) + ", lat="_s + nb::str( nb::float_( p.lat() ) ) + ")"_s;
} );
nb::class_<PointXY>( m, "PointXY" )
.def( nb::init<double,double>(), "x"_a, "y"_a )
.def_prop_ro( "x", nb::overload_cast<>( &PointXY::x, nb::const_ ) )
.def_prop_ro( "y", nb::overload_cast<>( &PointXY::y, nb::const_ ) )
.def( "__repr__", []( PointXY const& p ) {
return "_atlas4py.PointXY(x=" + std::to_string( p.x() ) + ", y=" + std::to_string( p.y() ) + ")";
return "atlas4py.PointXY(x="_s + nb::str( nb::float_( p.x() ) ) + ", y="_s + nb::str( nb::float_( p.y() ) ) + ")"_s;
} );

nb::class_<Projection>( m, "Projection" )
.def( "__repr__", []( Projection const& p ) {
return "_atlas4py.Projection("_s + nb::str( atlas4py::make_object( p.spec() ) ) + ")"_s;
return "<atlas4py.Projection("_s + nb::str( atlas4py::make_object( p.spec() ) ) + ")>"_s;
} );

nb::class_<Domain>( m, "Domain" )
Expand All @@ -252,9 +252,9 @@ NB_MODULE( _atlas4py, m ) {
.def_prop_ro( "units", &Domain::units )
.def( "__repr__", []( Domain const& d ) {
if (d) {
return nb::str("_atlas4py.Domain("_s + nb::str( atlas4py::make_object( d.spec() ) ) + ")"_s);
return nb::str("<atlas4py.Domain("_s + nb::str( atlas4py::make_object( d.spec() ) ) + ")>"_s);
}
return nb::str("_atlas4py.Domain()"_s);
return nb::str("<atlas4py.Domain()>"_s);
} );
nb::class_<RectangularDomain, Domain>( m, "RectangularDomain" )
.def( nb::init<const atlas::RectangularDomain::Interval&, const atlas::RectangularDomain::Interval&, const std::string&>(), "x_interval"_a, "y_interval"_a, "units"_a = "degrees" );
Expand All @@ -267,13 +267,13 @@ NB_MODULE( _atlas4py, m ) {
.def_prop_ro( "projection", &Grid::projection )
.def_prop_ro( "domain", &Grid::domain )
.def( "__repr__",
[]( Grid const& g ) { return "_atlas4py.Grid("_s + nb::str( atlas4py::make_object( g.spec() ) ) + ")"_s; } );
[]( Grid const& g ) { return "<atlas4py.Grid("_s + nb::str( atlas4py::make_object( g.spec() ) ) + ")>"_s; } );

nb::class_<grid::Spacing>( m, "Spacing" )
.def( "__len__", &grid::Spacing::size )
.def( "__getitem__", &grid::Spacing::operator[])
.def( "__repr__", []( grid::Spacing const& spacing ) {
return "_atlas4py.Spacing("_s + nb::str( atlas4py::make_object( spacing.spec() ) ) + ")"_s;
return "<atlas4py.Spacing("_s + nb::str( atlas4py::make_object( spacing.spec() ) ) + ")>"_s;
} );
nb::class_<grid::LinearSpacing, grid::Spacing>( m, "LinearSpacing" )
.def( nb::init<double,double,long,bool>(), "start"_a, "stop"_a, "N"_a, "endpoint_included"_a = true );
Expand Down Expand Up @@ -343,20 +343,50 @@ NB_MODULE( _atlas4py, m ) {
})
.def("__dlpack_device__", [](nb::handle /*self*/) {
return std::make_pair(nb::device::cpu::value, 0);
})
.def("__repr__", []( Field const& field ) -> std::string {
std::ostringstream oss;
Comment thread
wdeconinck marked this conversation as resolved.
oss << "<atlas4py.Field"
<< " name=" << field.name()
<< " shape=(";
for (size_t i = 0; i < field.shape().size(); ++i) {
if (i > 0) oss << ", ";
oss << field.shape()[i];
}
if (field.shape().size() == 1) oss << ",";
oss << ")"
Comment thread
Copilot marked this conversation as resolved.
<< " dtype=" << atlas4py::dtype::to_python_name(field.datatype())
<< ">";
return oss.str();
});

nb::class_<Mesh>( m, "Mesh" )
.def_prop_ro( "grid", &Mesh::grid )
.def_prop_ro( "projection", &Mesh::projection )
.def_prop_ro( "nodes", nb::overload_cast<>( &Mesh::nodes, nb::const_ ))
.def_prop_ro( "edges", nb::overload_cast<>( &Mesh::edges, nb::const_ ))
.def_prop_ro( "cells", nb::overload_cast<>( &Mesh::cells, nb::const_ ));
.def_prop_ro( "cells", nb::overload_cast<>( &Mesh::cells, nb::const_ ))
.def("__repr__", []( Mesh const& mesh ) -> std::string {
size_t halo = 0;
mesh.metadata().get("halo", halo);
std::ostringstream oss;
oss << "<atlas4py.Mesh"
<< " nb_nodes=" << mesh.nodes().size()
<< " nb_cells=" << mesh.cells().size()
<< " nb_edges=" << mesh.edges().size()
<< " halo=" << halo
<< ">";
return oss.str();
} );

nb::class_<StructuredMeshGenerator>( m, "StructuredMeshGenerator" )
// TODO in FunctionSpace below we expose config options, not the whole config object
.def( nb::init<const util::Config&>(), "config"_a )
.def( nb::init() )
.def( "generate", nb::overload_cast<Grid const&>( &StructuredMeshGenerator::generate, nb::const_ ) );
.def( "generate", nb::overload_cast<Grid const&>( &StructuredMeshGenerator::generate, nb::const_ ) )
.def( "__repr__", []( StructuredMeshGenerator const& smg ) -> std::string {
return "<atlas4py.StructuredMeshGenerator>";
} );

m.def( "build_edges", []( Mesh& mesh, const eckit::Configuration& config ) {
mesh::actions::build_edges( mesh, config);
Expand All @@ -381,7 +411,17 @@ NB_MODULE( _atlas4py, m ) {
.def_prop_ro( "rows", &mesh::IrregularConnectivity::rows )
.def( "cols", &mesh::IrregularConnectivity::cols, "row_idx"_a )
.def_prop_ro( "maxcols", &mesh::IrregularConnectivity::maxcols )
.def_prop_ro( "mincols", &mesh::IrregularConnectivity::mincols );
.def_prop_ro( "mincols", &mesh::IrregularConnectivity::mincols )
.def("__repr__", []( mesh::IrregularConnectivity const& c ) {
std::ostringstream oss;
if (c.rows() == 0) {
oss << "<atlas4py.IrregularConnectivity empty>";
}
else {
oss << "<atlas4py.IrregularConnectivity rows=" << c.rows() << " mincols=" << c.mincols() << " maxcols=" << c.maxcols() << ">";
}
return oss.str();
} );

nb::class_<mesh::BlockConnectivity>( m, "BlockConnectivity" )
.def( "__getitem__",
Expand All @@ -390,7 +430,17 @@ NB_MODULE( _atlas4py, m ) {
return c( row, col );
} )
.def_prop_ro( "rows", &mesh::BlockConnectivity::rows )
.def_prop_ro( "cols", &mesh::BlockConnectivity::cols );
.def_prop_ro( "cols", &mesh::BlockConnectivity::cols )
.def("__repr__", []( mesh::BlockConnectivity const& c ) {
std::ostringstream oss;
if (c.rows() == 0) {
oss << "<atlas4py.BlockConnectivity empty>";
}
else {
oss << "<atlas4py.BlockConnectivity rows=" << c.rows() << " cols=" << c.cols() << ">";
}
return oss.str();
} );

nb::class_<mesh::MultiBlockConnectivity>( m, "MultiBlockConnectivity" )
.def( "__getitem__",
Expand All @@ -408,15 +458,30 @@ NB_MODULE( _atlas4py, m ) {
.def_prop_ro( "maxcols", &mesh::MultiBlockConnectivity::maxcols )
.def_prop_ro( "mincols", &mesh::MultiBlockConnectivity::mincols )
.def_prop_ro( "blocks", &mesh::MultiBlockConnectivity::blocks )
.def( "block", nb::overload_cast<idx_t>( &mesh::MultiBlockConnectivity::block, nb::const_ ), nb::rv_policy::reference_internal );
.def( "block", nb::overload_cast<idx_t>( &mesh::MultiBlockConnectivity::block, nb::const_ ), nb::rv_policy::reference_internal )
.def("__repr__", []( mesh::MultiBlockConnectivity const& c ) {
std::ostringstream oss;
if (c.rows() == 0) {
oss << "<atlas4py.MultiBlockConnectivity empty>";
}
else {
oss << "<atlas4py.MultiBlockConnectivity blocks=" << c.blocks() << " rows=" << c.rows() << " mincols=" << c.mincols() << " maxcols=" << c.maxcols() << ">";
}
return oss.str();
} );

nb::class_<mesh::Nodes>( m, "Nodes" )
.def_prop_ro( "size", &mesh::Nodes::size )
.def_prop_ro( "edge_connectivity", nb::overload_cast<>( &mesh::Nodes::edge_connectivity, nb::const_ ) )
.def_prop_ro( "cell_connectivity", nb::overload_cast<>( &mesh::Nodes::cell_connectivity, nb::const_ ) )
.def_prop_ro( "lonlat", nb::overload_cast<>( &Mesh::Nodes::lonlat, nb::const_ ) )
.def_prop_ro( "lonlat", nb::overload_cast<>( &mesh::Nodes::lonlat, nb::const_ ) )
.def("field", []( mesh::Nodes const& n, std::string const& name ) { return n.field( name ); }, "name"_a, nb::rv_policy::reference_internal )
.def( "flags", []( mesh::Nodes const& n ) { return n.flags(); }, nb::rv_policy::reference_internal);
.def( "flags", []( mesh::Nodes const& n ) { return n.flags(); }, nb::rv_policy::reference_internal)
.def("__repr__", []( mesh::Nodes const& n ) {
std::ostringstream oss;
oss << "<atlas4py.Nodes size=" << n.size() << ">";
return oss.str();
} );

nb::class_<mesh::HybridElements>( m, "HybridElements" )
.def_prop_ro( "size", &mesh::HybridElements::size )
Expand All @@ -426,7 +491,12 @@ NB_MODULE( _atlas4py, m ) {
.def_prop_ro( "edge_connectivity", nb::overload_cast<>( &mesh::HybridElements::edge_connectivity, nb::const_ ) )
.def_prop_ro( "cell_connectivity", nb::overload_cast<>( &mesh::HybridElements::cell_connectivity, nb::const_ ) )
.def( "field", []( mesh::HybridElements const& he, std::string const& name ) { return he.field( name ); }, "name"_a, nb::rv_policy::reference_internal )
.def( "flags", []( mesh::HybridElements const& he ) { return he.flags(); }, nb::rv_policy::reference_internal );
.def( "flags", []( mesh::HybridElements const& he ) { return he.flags(); }, nb::rv_policy::reference_internal )
.def("__repr__", []( mesh::HybridElements const& he ) {
std::ostringstream oss;
oss << "<atlas4py.HybridElements size=" << he.size() << ">";
return oss.str();
} );

auto m_fs = m.def_submodule( "functionspace" );
nb::class_<FunctionSpace>( m_fs, "FunctionSpace" )
Expand All @@ -450,27 +520,59 @@ NB_MODULE( _atlas4py, m ) {
config = config | option::variables( *variables );
config = config | option::datatype( atlas4py::dtype::from_python_object( dtype ) );
return fs.createField( config );
}, "dtype"_a, "name"_a = std::nullopt, "levels"_a = std::nullopt, "variables"_a = std::nullopt );
}, "dtype"_a, "name"_a = std::nullopt, "levels"_a = std::nullopt, "variables"_a = std::nullopt )
.def("__repr__", []( FunctionSpace const& fs ) {
std::ostringstream oss;
oss << "<atlas4py.functionspace.FunctionSpace type=" << fs.type() << " size=" << fs.size() << ">";
return oss.str();
} );

nb::class_<functionspace::EdgeColumns, FunctionSpace>( m_fs, "EdgeColumns" )
.def("__init__", [](functionspace::EdgeColumns *t, const Mesh&m, int halo) { new (t) functionspace::EdgeColumns(m, util::Config()( "halo", halo )); }, "mesh"_a, "halo"_a = 0 )
.def_prop_ro( "nb_edges", &functionspace::EdgeColumns::nb_edges )
.def_prop_ro( "mesh", &functionspace::EdgeColumns::mesh )
.def_prop_ro( "edges", &functionspace::EdgeColumns::edges )
.def_prop_ro( "valid", &functionspace::EdgeColumns::valid );
.def_prop_ro( "valid", &functionspace::EdgeColumns::valid )
.def("__repr__", []( functionspace::EdgeColumns const& ec ) -> std::string {
std::ostringstream oss;
if ( !ec.valid() ) {
oss << "<atlas4py.functionspace.EdgeColumns invalid>";
} else {
oss << "<atlas4py.functionspace.EdgeColumns size=" << ec.size() << ">";
}
return oss.str();
} );

nb::class_<functionspace::NodeColumns, FunctionSpace>( m_fs, "NodeColumns" )
.def("__init__", [](functionspace::NodeColumns *t, const Mesh&m, int halo) { new (t) functionspace::NodeColumns(m, util::Config()( "halo", halo )); }, "mesh"_a, "halo"_a = 0 )
.def_prop_ro( "nb_nodes", &functionspace::NodeColumns::nb_nodes )
.def_prop_ro( "mesh", &functionspace::NodeColumns::mesh )
.def_prop_ro( "nodes", &functionspace::NodeColumns::nodes )
.def_prop_ro( "valid", &functionspace::NodeColumns::valid );
.def_prop_ro( "valid", &functionspace::NodeColumns::valid )
.def("__repr__", []( functionspace::NodeColumns const& nc ) -> std::string {
std::ostringstream oss;
if ( !nc.valid() ) {
oss << "<atlas4py.functionspace.NodeColumns invalid>";
} else {
oss << "<atlas4py.functionspace.NodeColumns size=" << nc.nb_nodes() << ">";
}
return oss.str();
} );
nb::class_<functionspace::CellColumns, FunctionSpace>( m_fs, "CellColumns" )
.def("__init__", [](functionspace::CellColumns *t, const Mesh&m, int halo) { new (t) functionspace::CellColumns(m, util::Config()( "halo", halo )); }, "mesh"_a, "halo"_a = 0 )
.def_prop_ro( "nb_cells", &functionspace::CellColumns::nb_cells )
.def_prop_ro( "mesh", &functionspace::CellColumns::mesh )
.def_prop_ro( "cells", &functionspace::CellColumns::cells )
.def_prop_ro( "valid", &functionspace::CellColumns::valid );
.def_prop_ro( "valid", &functionspace::CellColumns::valid )
.def("__repr__", []( functionspace::CellColumns const& cc ) -> std::string {
std::ostringstream oss;
if ( !cc.valid() ) {
oss << "<atlas4py.functionspace.CellColumns invalid>";
} else {
oss << "<atlas4py.functionspace.CellColumns size=" << cc.nb_cells() << ">";
}
return oss.str();
} );

nb::class_<util::Metadata, eckit::LocalConfiguration>( m, "Metadata" )
.def( "__repr__", []( util::Metadata const& metadata ) {
Expand All @@ -495,6 +597,9 @@ NB_MODULE( _atlas4py, m ) {
topology.def_static( "check", &mesh::Nodes::Topology::check );
topology.def_static( "check_all", &mesh::Nodes::Topology::check_all );
topology.def_static( "check_any", &mesh::Nodes::Topology::check_any );
topology.def("__repr__", []( mesh::Nodes::Topology const& topology ) {
return "<atlas4py.Topology>";
} );

nb::class_<output::Gmsh>( m, "Gmsh" )
.def( nb::init<std::string const&>(), "path"_a )
Expand All @@ -503,5 +608,5 @@ NB_MODULE( _atlas4py, m ) {
.def( "write", []( output::Gmsh& gmsh, Mesh const& mesh ) { gmsh.write( mesh ); }, "mesh"_a )
.def( "write", []( output::Gmsh& gmsh, Field const& field ) { gmsh.write( field ); }, "field"_a )
.def( "write", []( output::Gmsh& gmsh, Field const& field, FunctionSpace const& fs ) { gmsh.write( field, fs ); }, "field"_a, "functionspace"_a )
.def("__repr__", []( output::Gmsh const& gmsh ) { return "_atlas4py.output.Gmsh()"; } );
.def("__repr__", []( output::Gmsh const& gmsh ) { return "<atlas4py.Gmsh>"; } );
}
67 changes: 67 additions & 0 deletions tests/test_bindings.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,21 @@ def test_version():
assert isinstance(atlas4py.__version__, str)


@pytest.mark.parametrize(
("point", "coordinates"),
[
(atlas4py.PointLonLat(12.345678901234, -3.456789012345), ("lon", "lat")),
(atlas4py.PointXY(-2.345678901234, 8.456789012345), ("x", "y")),
],
)
def test_point_repr_roundtrip(point, coordinates):
Comment thread
wdeconinck marked this conversation as resolved.
reconstructed = eval(repr(point))

assert type(reconstructed) is type(point)
for coordinate in coordinates:
assert getattr(reconstructed, coordinate) == getattr(point, coordinate)


def test_grid_generation(structured_grid):
assert structured_grid.domain.type == "rectangular"
assert structured_grid.regular == True
Expand Down Expand Up @@ -117,10 +132,55 @@ def test_mesh_connectivity(structured_mesh):
assert block[0, 3] == 21


def test_mesh_and_connectivity_repr_branches():
grid = atlas4py.StructuredGrid(
x_spacing=atlas4py.LinearSpacing(-1, 1, 4),
y_spacing=atlas4py.LinearSpacing(-1, 1, 3),
)
generator = atlas4py.StructuredMeshGenerator()
mesh = generator.generate(grid)

assert repr(generator) == "<atlas4py.StructuredMeshGenerator>"
assert repr(mesh) == "<atlas4py.Mesh nb_nodes=12 nb_cells=6 nb_edges=0 halo=0>"
assert repr(mesh.nodes) == "<atlas4py.Nodes size=12>"
assert repr(mesh.cells) == "<atlas4py.HybridElements size=6>"
assert repr(mesh.nodes.edge_connectivity) == "<atlas4py.IrregularConnectivity empty>"
assert repr(mesh.nodes.cell_connectivity) == "<atlas4py.IrregularConnectivity empty>"
assert repr(mesh.cells.edge_connectivity) == "<atlas4py.MultiBlockConnectivity empty>"
assert repr(mesh.cells.cell_connectivity) == "<atlas4py.MultiBlockConnectivity empty>"
assert repr(mesh.cells.node_connectivity) == (
"<atlas4py.MultiBlockConnectivity blocks=2 rows=6 mincols=3 maxcols=4>"
)
assert repr(mesh.cells.node_connectivity.block(0)) == (
"<atlas4py.BlockConnectivity rows=6 cols=4>"
)

atlas4py.build_edges(mesh)
atlas4py.build_node_to_edge_connectivity(mesh)

assert repr(mesh) == "<atlas4py.Mesh nb_nodes=12 nb_cells=6 nb_edges=17 halo=0>"
assert repr(mesh.nodes.edge_connectivity) == (
"<atlas4py.IrregularConnectivity rows=12 mincols=2 maxcols=4>"
)


def test_function_space_generation(structured_function_space):
assert structured_function_space.nb_cells == 380


def test_function_space_repr(structured_mesh):
edge_columns = atlas4py.functionspace.EdgeColumns(structured_mesh)
node_columns = atlas4py.functionspace.NodeColumns(structured_mesh)
cell_columns = atlas4py.functionspace.CellColumns(structured_mesh)

assert repr(edge_columns) == "<atlas4py.functionspace.EdgeColumns size=799>"
assert repr(node_columns) == "<atlas4py.functionspace.NodeColumns size=420>"
assert repr(cell_columns) == "<atlas4py.functionspace.CellColumns size=380>"
assert atlas4py.functionspace.FunctionSpace.__repr__(cell_columns) == (
"<atlas4py.functionspace.FunctionSpace type=CellColumns size=380>"
)


def test_field_generation(structured_in_and_out_fields):
in_f, out_f = structured_in_and_out_fields
assert in_f.rank == 2
Expand All @@ -136,6 +196,8 @@ def test_field_generation(structured_in_and_out_fields):

assert np.allclose(in_view + 1, out_view)

assert repr(in_f) == "<atlas4py.Field name=my_in_field shape=(380, 1) dtype=float64>"


def test_metadata_mapping_protocol(structured_in_and_out_fields):
field, _ = structured_in_and_out_fields
Expand All @@ -151,6 +213,10 @@ def test_metadata_mapping_protocol(structured_in_and_out_fields):
assert len(metadata) == len(values)
assert "source" in metadata
assert metadata["source"] == "test"
assert repr(metadata) == (
"atlas4py.Metadata({'name': 'my_in_field', 'levels': 1, 'variables': 0, "
"'global': False, 'source': 'test', 'level': 1})"
)


def test_field_array_accepts_matching_dtype_and_false_copy(structured_in_and_out_fields):
Expand All @@ -171,6 +237,7 @@ def test_gmsh_output(structured_mesh):

output_file = "test_output.msh"
with atlas4py.Gmsh(path=output_file) as gmsh:
assert repr(gmsh) == "<atlas4py.Gmsh>"
gmsh.write(structured_mesh)

# Check that the file was created and has content
Expand Down
Loading