From ba69610083bffc3a1487ef1d804af336b80d7c05 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Fri, 27 May 2022 16:01:34 +0100 Subject: [PATCH 01/20] Testing and improving. Works with atlas 0.29.0 --- examples/example1.py | 25 +++++++ setup.py | 2 +- src/atlas4py/CMakeLists.txt | 35 +++++----- src/atlas4py/__init__.py | 5 ++ src/atlas4py/_atlas4py.cpp | 130 ++++++++++++++++++++++++++---------- 5 files changed, 146 insertions(+), 51 deletions(-) create mode 100755 examples/example1.py diff --git a/examples/example1.py b/examples/example1.py new file mode 100755 index 0000000..1d1eb3d --- /dev/null +++ b/examples/example1.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python3 +import atlas4py as atlas +import math + +atlas.initialize() # Required + +mesh = atlas.Mesh( atlas.Grid("H16") ) + +fs_nodes = atlas.functionspace.NodeColumns(mesh) + +field = fs_nodes.create_field(name="myfield") + +lonlat = atlas.make_view(fs_nodes.lonlat) + +view = atlas.make_view(field) +for n in range(field.shape[0]): + lat = lonlat[n,1] * math.pi / 180. + view[n] = math.cos(4.*lat) + +gmsh = atlas.Gmsh("out.msh", coordinates='xyz') +gmsh.write(mesh) +gmsh.write(field) + +atlas.finalize() + diff --git a/setup.py b/setup.py index eeb22a3..2a95202 100644 --- a/setup.py +++ b/setup.py @@ -121,7 +121,7 @@ def build_extension(self, ext): subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=self.build_temp) -PACKAGE_VERSION = "0.26.0.dev14" +PACKAGE_VERSION = "0.29.0.dev14" # Meaning of the version scheme "{major}.{minor}.{patch}.dev{dev}": # - {major}.{minor}.{patch} => version of the atlas C++ library (hardcoded in 'setup.py') # - {dev} => version of the Python bindings as the commit number in 'master' diff --git a/src/atlas4py/CMakeLists.txt b/src/atlas4py/CMakeLists.txt index 15d8134..e96b9a3 100644 --- a/src/atlas4py/CMakeLists.txt +++ b/src/atlas4py/CMakeLists.txt @@ -2,19 +2,10 @@ cmake_minimum_required(VERSION ${ATLAS4PY_CMAKE_MINIMUM_REQUIRED_VERSION}) project(atlas4py LANGUAGES CXX) +cmake_policy(SET CMP0074 NEW) set(CMAKE_CXX_STANDARD 17) include(FetchContent) -macro(find_ecbuild) - if(NOT ecbuild_FOUND) - FetchContent_Declare( - ecbuild - GIT_REPOSITORY https://github.com/ecmwf/ecbuild.git - GIT_TAG ${ATLAS4PY_ECBUILD_VERSION} - ) - FetchContent_MakeAvailable(ecbuild) - endif() -endmacro() function(copy_target ct_tgt) add_custom_target(copy_${ct_tgt} ALL @@ -24,21 +15,35 @@ function(copy_target ct_tgt) endfunction() find_package(atlas QUIET PATHS $ENV{ATLAS_INSTALL_DIR}) +if( atlas_FOUND ) + message( "Found atlas: ${atlas_DIR} (found version \"${atlas_VERSION}\")" ) +endif() + if(NOT atlas_FOUND) - find_package(eckit QUIET) + find_package(ecbuild) + if(NOT ecbuild_FOUND) + FetchContent_Declare( + ecbuild + GIT_REPOSITORY https://github.com/ecmwf/ecbuild.git + GIT_TAG ${ATLAS4PY_ECBUILD_VERSION} + ) + FetchContent_MakeAvailable(ecbuild) + endif() + find_package(eckit) + if( eckit_FOUND ) + message( "Found eckit: ${eckit_DIR} (found version \"${eckit_VERSION}\")" ) + endif() + if(NOT eckit_FOUND) - find_ecbuild() FetchContent_Declare( eckit GIT_REPOSITORY https://github.com/ecmwf/eckit.git GIT_TAG ${ATLAS4PY_ECKIT_VERSION} ) FetchContent_MakeAvailable(eckit) - set(_atlas4py_built_eckit ON) endif() - find_ecbuild() FetchContent_Declare( atlas GIT_REPOSITORY https://github.com/ecmwf/atlas.git @@ -70,4 +75,4 @@ FetchContent_Declare( FetchContent_MakeAvailable(pybind11) pybind11_add_module(_atlas4py _atlas4py.cpp) -target_link_libraries(_atlas4py PUBLIC atlas eckit) +target_link_libraries(_atlas4py PUBLIC atlas) diff --git a/src/atlas4py/__init__.py b/src/atlas4py/__init__.py index ffa06e6..5212c44 100644 --- a/src/atlas4py/__init__.py +++ b/src/atlas4py/__init__.py @@ -6,3 +6,8 @@ from ._version import __version__ from ._atlas4py import * + +def make_view(field): + import numpy as np + return np.array(field, copy=False) + diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index 91f6f92..4633507 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -16,9 +16,12 @@ #include "atlas/meshgenerator.h" #include "atlas/option/Options.h" #include "atlas/output/Gmsh.h" +#include "atlas/library.h" + #include "eckit/value/Value.h" #include "eckit/config/Configuration.h" + namespace py = ::pybind11; using namespace atlas; using namespace pybind11::literals; @@ -36,6 +39,48 @@ struct type_caster : public type_caster( value ) ) { + config.set(key,value.cast()); + } + else if ( py::isinstance( value ) ) { + config.set(key,value.cast()); + } + else if ( py::isinstance( value ) ) { + config.set(key,value.cast()); + } + else if ( py::isinstance( value ) ) { + config.set(key, value.cast()); + } + else { + throw std::out_of_range( "type of value unsupported" ); + } +} + +util::Config to_config( py::kwargs kwargs ) { + util::Config config; + for( const auto& pair : kwargs ) { + const auto key = pair.first.cast(); + const auto& value = pair.second; + config_set(config, key, value); + } + return config; +} + py::object toPyObject( eckit::Value const& v ) { if ( v.isBool() ) return py::bool_( v.as() ); @@ -195,34 +240,30 @@ PYBIND11_MODULE( _atlas4py, m ) { // TODO This is a duplicate of metadata below (because same base class) py::class_( m, "Config" ) .def( py::init() ) + .def( py::init( []( py::kwargs kwargs) { + return to_config(kwargs); + } ) ) .def( "__setitem__", []( util::Config& config, std::string const& key, py::object value ) { - if ( py::isinstance( value ) ) - config.set( key, value.cast() ); - else if ( py::isinstance( value ) ) - config.set( key, value.cast() ); - else if ( py::isinstance( value ) ) - config.set( key, value.cast() ); - else if ( py::isinstance( value ) ) - config.set( key, value.cast() ); - else - throw std::out_of_range( "type of value unsupported" ); + config_set(config,key,value); } ) .def( "__getitem__", []( util::Config& config, std::string const& key ) -> py::object { if ( !config.has( key ) ) throw std::out_of_range( "key <" + key + "> could not be found" ); - - // TODO: We have to query metadata.get() even though this should - // not be done (see comment in Config::get). We cannot - // avoid this right now because otherwise we cannot query - // the type of the underlying data. return toPyObject( config.get().element( key ) ); } ) .def( "__repr__", []( util::Config const& config ) { return "_atlas4py.Config("_s + py::str( toPyObject( config.get() ) ) + ")"_s; } ); + py::class_( m, "MeshGenerator" ) + // TODO in FunctionSpace below we expose config options, not the whole config object + .def( py::init( []( util::Config const& config ) { return MeshGenerator( config ); } ) ) + .def( py::init( []( const std::string& type ) { return MeshGenerator(type); } ) ) + .def( py::init( [](){ return MeshGenerator(); } ) ) + .def( "generate", py::overload_cast( &MeshGenerator::generate, py::const_ ) ); + py::class_( m, "StructuredMeshGenerator" ) // TODO in FunctionSpace below we expose config options, not the whole config object .def( py::init( []( util::Config const& config ) { return StructuredMeshGenerator( config ); } ) ) @@ -230,6 +271,7 @@ PYBIND11_MODULE( _atlas4py, m ) { .def( "generate", py::overload_cast( &StructuredMeshGenerator::generate, py::const_ ) ); py::class_( m, "Mesh" ) + .def( py::init( []( const Grid& grid ) { return Mesh(grid); } ) ) .def_property_readonly( "grid", &Mesh::grid ) .def_property_readonly( "projection", &Mesh::projection ) .def_property( "nodes", py::overload_cast<>( &Mesh::nodes, py::const_ ), py::overload_cast<>( &Mesh::nodes ) ) @@ -317,6 +359,19 @@ PYBIND11_MODULE( _atlas4py, m ) { py::return_value_policy::reference_internal ); + + auto m_library = m.def_submodule( "library" ); + + m_library.def("initialize", []() {initialise_sys_argv();}); + m_library.def("finalize", [](){atlas::finalise(); } ); + m_library.def("initialise", []() {initialise_sys_argv();}); + m_library.def("finalise", [](){atlas::finalise(); } ); + m_library.attr("version") = atlas::Library::instance().version(); + m.def("initialize", []() {initialise_sys_argv();}); + m.def("finalize", []() {atlas::finalize();}); + m.def("initialise", []() {initialise_sys_argv();}); + m.def("finalise", []() {atlas::finalize();}); + auto m_fs = m.def_submodule( "functionspace" ); py::class_( m_fs, "FunctionSpace" ) .def_property_readonly( "size", &FunctionSpace::size ) @@ -324,42 +379,46 @@ PYBIND11_MODULE( _atlas4py, m ) { .def( "create_field", []( FunctionSpace const& fs, std::optional const& name, std::optional levels, - std::optional variables, py::object dtype ) { + std::optional variables, std::optional dtype ) { util::Config config; if ( name ) - config = config | option::name( *name ); + config.set(option::name( *name )); // TODO what does it mean in atlas if levels is not set? if ( levels ) - config = config | option::levels( *levels ); + config.set(option::levels( *levels )); if ( variables ) - config = config | option::variables( *variables ); - config = config | option::datatype( pybindToAtlas( py::dtype::from_args( dtype ) ) ); + config.set(option::variables( *variables )); + if ( dtype ) + config.set( option::datatype( pybindToAtlas( py::dtype::from_args( *dtype ) ) )); + else + config.set( option::datatypeT() ); return fs.createField( config ); }, - "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, "dtype"_a = std::nullopt) + .def_property_readonly("lonlat", &FunctionSpace::lonlat ); py::class_( m_fs, "EdgeColumns" ) - .def( py::init( []( Mesh const& m, int halo ) { - return functionspace::EdgeColumns( m, util::Config()( "halo", halo ) ); + .def( py::init( []( Mesh const& m, int halo, int levels ) { + return functionspace::EdgeColumns( m, util::Config()( "halo", halo )("levels", levels) ); } ), - "mesh"_a, "halo"_a = 0 ) + "mesh"_a, "halo"_a = 0, "levels"_a = 0 ) .def_property_readonly( "nb_edges", &functionspace::EdgeColumns::nb_edges ) .def_property_readonly( "mesh", &functionspace::EdgeColumns::mesh ) .def_property_readonly( "edges", &functionspace::EdgeColumns::edges ) .def_property_readonly( "valid", &functionspace::EdgeColumns::valid ); py::class_( m_fs, "NodeColumns" ) - .def( py::init( []( Mesh const& m, int halo ) { - return functionspace::NodeColumns( m, util::Config()( "halo", halo ) ); + .def( py::init( []( Mesh const& m, int halo, int levels ) { + return functionspace::NodeColumns( m, util::Config()( "halo", halo )("levels",levels) ); } ), - "mesh"_a, "halo"_a = 0 ) + "mesh"_a, "halo"_a = 0, "levels"_a = 0 ) .def_property_readonly( "nb_nodes", &functionspace::NodeColumns::nb_nodes ) .def_property_readonly( "mesh", &functionspace::NodeColumns::mesh ) .def_property_readonly( "nodes", &functionspace::NodeColumns::nodes ) .def_property_readonly( "valid", &functionspace::NodeColumns::valid ); py::class_( m_fs, "CellColumns" ) - .def( py::init( []( Mesh const& m, int halo ) { - return functionspace::CellColumns( m, util::Config()( "halo", halo ) ); + .def( py::init( []( Mesh const& m, int halo, int levels ) { + return functionspace::CellColumns( m, util::Config()( "halo", halo )("levels",levels) ); } ), - "mesh"_a, "halo"_a = 0 ) + "mesh"_a, "halo"_a = 0, "levels"_a = 0 ) .def_property_readonly( "nb_cells", &functionspace::CellColumns::nb_cells ) .def_property_readonly( "mesh", &functionspace::CellColumns::mesh ) .def_property_readonly( "cells", &functionspace::CellColumns::cells ) @@ -384,11 +443,6 @@ PYBIND11_MODULE( _atlas4py, m ) { []( util::Metadata& metadata, std::string const& key ) -> py::object { if ( !metadata.has( key ) ) throw std::out_of_range( "key <" + key + "> could not be found" ); - - // TODO: We have to query metadata.get() even though this should - // not be done (see comment in Config::get). We cannot - // avoid this right now because otherwise we cannot query - // the type of the underlying data. return toPyObject( metadata.get().element( key ) ); } ) .def( "__repr__", []( util::Metadata const& metadata ) { @@ -433,6 +487,12 @@ PYBIND11_MODULE( _atlas4py, m ) { py::class_( m, "Gmsh" ) .def( py::init( []( std::string const& path ) { return output::Gmsh{ path }; } ), "path"_a ) + .def( py::init( []( std::string const& path, eckit::Configuration const& config, py::kwargs kwargs ) { + util::Config cfg = util::Config(config); + cfg.set(to_config(kwargs)); + return output::Gmsh(path,cfg); + }), "path"_a, "config"_a ) + .def( py::init( []( std::string const& path, py::kwargs kwargs ) {return output::Gmsh{ path, to_config(kwargs) }; } ), "path"_a ) .def( "__enter__", []( output::Gmsh& gmsh ) { return gmsh; } ) .def( "__exit__", []( output::Gmsh& gmsh, py::object exc_type, py::object exc_val, py::object exc_tb ) { gmsh.reset( nullptr ); } ) From ddd8306a7d568d8071298b5fa3e3971e4ab37cfb Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Mon, 6 Jun 2022 10:24:31 +0100 Subject: [PATCH 02/20] Allow building as standalone for quick compilation testing --- src/atlas4py/CMakeLists.txt | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/atlas4py/CMakeLists.txt b/src/atlas4py/CMakeLists.txt index e96b9a3..3e5c71d 100644 --- a/src/atlas4py/CMakeLists.txt +++ b/src/atlas4py/CMakeLists.txt @@ -1,3 +1,8 @@ + +if( NOT ATLAS4PY_CMAKE_MINIMUM_REQUIRED_VERSION) + set(ATLAS4PY_CMAKE_MINIMUM_REQUIRED_VERSION 3.14) +endif() + cmake_minimum_required(VERSION ${ATLAS4PY_CMAKE_MINIMUM_REQUIRED_VERSION}) project(atlas4py LANGUAGES CXX) @@ -67,12 +72,15 @@ if(NOT atlas_FOUND) endif() endif() -FetchContent_Declare( - pybind11 - GIT_REPOSITORY https://github.com/pybind/pybind11.git - GIT_TAG v${ATLAS4PY_PYBIND11_VERSION} -) -FetchContent_MakeAvailable(pybind11) +find_package( pybind11 ) +if( NOT pybind11_FOUND ) + FetchContent_Declare( + pybind11 + GIT_REPOSITORY https://github.com/pybind/pybind11.git + GIT_TAG v${ATLAS4PY_PYBIND11_VERSION} + ) + FetchContent_MakeAvailable(pybind11) +endif() pybind11_add_module(_atlas4py _atlas4py.cpp) target_link_libraries(_atlas4py PUBLIC atlas) From 3670676ddddf66bc77948f4113525d834c90d2a5 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Mon, 6 Jun 2022 10:26:06 +0100 Subject: [PATCH 03/20] Further extend with Trans example --- examples/example1.py | 2 +- examples/example2.py | 103 ++++++++++++++++++++++++++++++++ src/atlas4py/_atlas4py.cpp | 118 ++++++++++++++++++++++++++----------- 3 files changed, 187 insertions(+), 36 deletions(-) create mode 100755 examples/example2.py diff --git a/examples/example1.py b/examples/example1.py index 1d1eb3d..4971c35 100755 --- a/examples/example1.py +++ b/examples/example1.py @@ -17,7 +17,7 @@ lat = lonlat[n,1] * math.pi / 180. view[n] = math.cos(4.*lat) -gmsh = atlas.Gmsh("out.msh", coordinates='xyz') +gmsh = atlas.Gmsh("example1.msh", coordinates='xyz') gmsh.write(mesh) gmsh.write(field) diff --git a/examples/example2.py b/examples/example2.py new file mode 100755 index 0000000..a284e69 --- /dev/null +++ b/examples/example2.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 + +# Example program that initializes a spectral field with a +# spherical harmonic. Then the inverse transform is performed +# Then the result can be visualized with Gmsh: +# gmsh example2.msh + +# ------------------------------------------------------------------------ +# +# Limitations with Trans.backend('ectrans') +# ----------------------------------------- +# 1) Only implemented for global Gaussian grids +# +# ------------------------------------------------------------------------ +# +# Limitations with Trans.backend('local') +# --------------------------------------- +# 1) Only `invtrans` is implemented, and only MPI-serial +# +# ------------------------------------------------------------------------ +# +# Known issues with Trans.backend('local'), to be fixed +# ----------------------------------------------------- +# 1) We need to initialize trans with +# trans = atlas.Trans(grid, truncation) +# instead of +# trans = atlas.Trans(fs_gp, fs_sp) +# 2) Using atlas Field API, only level=1 is supported. +# 3) The Healpix grid is identified as unstructured +# because the grid wrongly gets cropped to a "west", and its +# structure gets lost +# +# ------------------------------------------------------------------------ + +import atlas4py as atlas + +# ------------------------------------------------------------------------ +# Functions +def set_spherical_harmonic(field_sp, n, m): + """ Set field_sp to a spherical harmonic """ + nflds = field_sp.shape[1] + spectral_functionspace = atlas.functionspace.Spectral(field_sp.functionspace) + if not spectral_functionspace: + raise TypeError("'field_sp.functionspace' could not be cast to " + + "'atlas.functionspace.Spectral'. It is of type " + + "'atlas.functionspace."+field_sp.functionspace.type+"'.") + + sp = atlas.make_view(field_sp) + sp[:,:] = 0. + def spherical_harmonic(re,im,_n,_m): + for jfld in range(nflds): + if n*(jfld+1) == _n and m*(jfld+1) == _m: + sp[re,jfld] = 1. + sp[im,jfld] = 1. + + spectral_functionspace.parallel_for( spherical_harmonic ) +# ------------------------------------------------------------------------ + + +# ------------------------------------------------------------------------ +# Start here +# ------------------------------------------------------------------------ + +atlas.initialize() # Required + +grid = atlas.Grid("O128") +levels = 10 +truncation = 255 +atlas.Trans.backend("ectrans") +partitioner = atlas.Partitioner("ectrans") + # 'ectrans' partitioner needs to be used for 'ectrans' Trans backend only. + # It is essentially nearly identical to the equal_regions partitioner + +if not atlas.GaussianGrid(grid) or not atlas.Trans.has_backend("ectrans"): + # Fallback to "local" backend + atlas.Trans.backend("local") + levels = 1 + partitioner = atlas.Partitioner("equal_regions") + +# Create function spaces +fs_sp = atlas.functionspace.Spectral(truncation) +fs_gp = atlas.functionspace.StructuredColumns(grid, partitioner) + +# Create fields (already distributed) +field_sp = fs_sp.create_field(name="sp", levels=levels) +field_gp = fs_gp.create_field(name="gp", levels=levels) + +# Initial condition for spectral field +set_spherical_harmonic(field_sp, n=4, m=2) + +# Spectral transform +trans = atlas.Trans(grid, truncation) +trans.invtrans(field_sp, field_gp) + +# Visualisation of gridpoint field +meshgenerator = atlas.MeshGenerator(type='structured') +mesh = meshgenerator.generate(grid) +gmsh = atlas.Gmsh("example2.msh", coordinates='xy') +gmsh.write(mesh) +gmsh.write(field_gp) + +atlas.finalize() +# ------------------------------------------------------------------------ diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index 4633507..8248e47 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -17,6 +17,7 @@ #include "atlas/option/Options.h" #include "atlas/output/Gmsh.h" #include "atlas/library.h" +#include "atlas/trans/Trans.h" #include "eckit/value/Value.h" #include "eckit/config/Configuration.h" @@ -39,18 +40,27 @@ struct type_caster : public type_caster( m, "PointLonLat" ) .def( py::init( []( double lon, double lat ) { return PointLonLat( { lon, lat } ); @@ -234,6 +256,14 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "regular", &StructuredGrid::regular ) .def_property_readonly( "periodic", &StructuredGrid::periodic ); + py::class_( m, "GaussianGrid" ) + .def( py::init( []( Grid const& g ) { + return GaussianGrid{ g }; + } ), + "grid"_a ) + .def_property_readonly( "valid", &GaussianGrid::valid ) + .def("__bool__", &GaussianGrid::valid); + py::class_( m, "eckit.Configuration" ); py::class_( m, "eckit.LocalConfiguration" ); @@ -257,8 +287,14 @@ PYBIND11_MODULE( _atlas4py, m ) { return "_atlas4py.Config("_s + py::str( toPyObject( config.get() ) ) + ")"_s; } ); + py::class_( m, "Partitioner" ) + .def( py::init( []( py::kwargs kwargs ) { return grid::Partitioner( to_config(kwargs)); } ) ) + .def( py::init( []( util::Config const& config ) { return grid::Partitioner( config ); } ) ) + .def( py::init( []( const std::string& type ) { return grid::Partitioner(type); } ) ) + .def( py::init( [](){ return grid::Partitioner(); } ) ); + py::class_( m, "MeshGenerator" ) - // TODO in FunctionSpace below we expose config options, not the whole config object + .def( py::init( []( py::kwargs kwargs ) { return MeshGenerator( to_config(kwargs)); } ) ) .def( py::init( []( util::Config const& config ) { return MeshGenerator( config ); } ) ) .def( py::init( []( const std::string& type ) { return MeshGenerator(type); } ) ) .def( py::init( [](){ return MeshGenerator(); } ) ) @@ -358,43 +394,21 @@ PYBIND11_MODULE( _atlas4py, m ) { .def( "flags", []( mesh::HybridElements const& he ) { return he.flags(); }, py::return_value_policy::reference_internal ); - - - auto m_library = m.def_submodule( "library" ); - - m_library.def("initialize", []() {initialise_sys_argv();}); - m_library.def("finalize", [](){atlas::finalise(); } ); - m_library.def("initialise", []() {initialise_sys_argv();}); - m_library.def("finalise", [](){atlas::finalise(); } ); - m_library.attr("version") = atlas::Library::instance().version(); - m.def("initialize", []() {initialise_sys_argv();}); - m.def("finalize", []() {atlas::finalize();}); - m.def("initialise", []() {initialise_sys_argv();}); - m.def("finalise", []() {atlas::finalize();}); - auto m_fs = m.def_submodule( "functionspace" ); py::class_( m_fs, "FunctionSpace" ) .def_property_readonly( "size", &FunctionSpace::size ) .def_property_readonly( "type", &FunctionSpace::type ) .def( "create_field", - []( FunctionSpace const& fs, std::optional const& name, std::optional levels, - std::optional variables, std::optional dtype ) { - util::Config config; - if ( name ) - config.set(option::name( *name )); - // TODO what does it mean in atlas if levels is not set? - if ( levels ) - config.set(option::levels( *levels )); - if ( variables ) - config.set(option::variables( *variables )); + []( FunctionSpace const& fs, std::optional dtype, py::kwargs kwargs ) { + util::Config config = to_config(kwargs); if ( dtype ) config.set( option::datatype( pybindToAtlas( py::dtype::from_args( *dtype ) ) )); else config.set( option::datatypeT() ); return fs.createField( config ); }, - "name"_a = std::nullopt, "levels"_a = std::nullopt, "variables"_a = std::nullopt, "dtype"_a = std::nullopt) + "dtype"_a = std::nullopt) .def_property_readonly("lonlat", &FunctionSpace::lonlat ); py::class_( m_fs, "EdgeColumns" ) .def( py::init( []( Mesh const& m, int halo, int levels ) { @@ -424,6 +438,26 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "cells", &functionspace::CellColumns::cells ) .def_property_readonly( "valid", &functionspace::CellColumns::valid ); + py::class_( m_fs, "Spectral" ) + .def( py::init( []( FunctionSpace fs ) { return functionspace::Spectral{fs}; } ) ) + .def( py::init( []( int truncation ) { return functionspace::Spectral( truncation ); } ), "truncation"_a ) + .def_property_readonly( "nb_spectral_coefficients", &functionspace::Spectral::nb_spectral_coefficients ) + .def_property_readonly( "nb_spectral_coefficients_global", &functionspace::Spectral::nb_spectral_coefficients_global ) + .def_property_readonly( "truncation", &functionspace::Spectral::truncation ) + .def_property_readonly( "valid", &functionspace::Spectral::valid ) + .def( "__bool__", [](const functionspace::Spectral& self) { return self.valid(); } ) + .def("parallel_for", []( const functionspace::Spectral& self, const py::function &f) { + self.parallel_for>(f);}); + + + py::class_( m_fs, "StructuredColumns" ) + .def( py::init( []( Grid const& g, grid::Partitioner const& p, py::kwargs kwargs ) { + return functionspace::StructuredColumns( g, p, to_config(kwargs) ); } ), + "grid"_a, "partitioner"_a ) + .def( py::init( []( Grid const& g, py::kwargs kwargs ) { return functionspace::StructuredColumns( g, to_config(kwargs) ); } ), "grid"_a ) + .def_property_readonly( "grid", &functionspace::StructuredColumns::grid ) + .def_property_readonly( "valid", &functionspace::StructuredColumns::valid ); + py::class_( m, "Metadata" ) .def_property_readonly( "keys", &util::Metadata::keys ) .def( "__setitem__", @@ -458,6 +492,7 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "datatype", []( Field& f ) { return atlasToPybind( f.datatype() ); } ) .def_property( "metadata", py::overload_cast<>( &Field::metadata, py::const_ ), py::overload_cast<>( &Field::metadata ) ) + .def_property_readonly( "functionspace", py::overload_cast<>( &Field::functionspace, py::const_ ) ) .def_buffer( []( Field& f ) { auto strides = f.strides(); std::transform( strides.begin(), strides.end(), strides.begin(), @@ -503,4 +538,17 @@ PYBIND11_MODULE( _atlas4py, m ) { .def( "write", []( output::Gmsh& gmsh, Field const& field, FunctionSpace const& fs ) { gmsh.write( field, fs ); }, "field"_a, "functionspace"_a ); + + + py::class_( m, "Trans" ) + .def( py::init( [](const FunctionSpace& gp, const FunctionSpace& sp){ return trans::Trans(gp,sp);} ), "gp"_a, "sp"_a ) + .def( py::init( [](const Grid& grid, int truncation){ return trans::Trans(grid,truncation);} ), "grid"_a, "truncation"_a ) + .def( "dirtrans", []( trans::Trans& trans, const Field& gpfield, Field& spfield) { trans.dirtrans(gpfield,spfield);} ) + .def( "invtrans", []( trans::Trans& trans, const Field& spfield, Field& gpfield) { trans.invtrans(spfield,gpfield);} ) + .def_property_readonly( "truncation", &trans::Trans::truncation ) + .def_property_readonly( "nb_spectral_coefficients", &trans::Trans::spectralCoefficients ) + .def_static("backend", [] ( const std::string& backend ){ trans::Trans::backend(backend); } ) + .def_static("has_backend", [] ( const std::string& backend ){ return trans::Trans::hasBackend(backend); } ); + + } From 149ed485718fd5eddbf1ab640498f58a4d26eb97 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Tue, 7 Jun 2022 16:25:48 +0100 Subject: [PATCH 04/20] Tweaks --- src/atlas4py/_atlas4py.cpp | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index 8248e47..f6971af 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -373,6 +373,7 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "cell_connectivity", py::overload_cast<>( &mesh::Nodes::cell_connectivity, py::const_ ) ) .def_property_readonly( "lonlat", py::overload_cast<>( &Mesh::Nodes::lonlat, py::const_ ) ) + .def_property_readonly( "xy", py::overload_cast<>( &Mesh::Nodes::xy, py::const_ ) ) .def("field", []( mesh::Nodes const& n, std::string const& name ) { return n.field( name ); }, "name"_a, py::return_value_policy::reference_internal ) .def( "flags", []( mesh::Nodes const& n ) { return n.flags(); }, @@ -420,6 +421,7 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "edges", &functionspace::EdgeColumns::edges ) .def_property_readonly( "valid", &functionspace::EdgeColumns::valid ); py::class_( m_fs, "NodeColumns" ) + .def( py::init( []( FunctionSpace fs ) { return functionspace::NodeColumns{fs}; } ) ) .def( py::init( []( Mesh const& m, int halo, int levels ) { return functionspace::NodeColumns( m, util::Config()( "halo", halo )("levels",levels) ); } ), @@ -427,7 +429,9 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "nb_nodes", &functionspace::NodeColumns::nb_nodes ) .def_property_readonly( "mesh", &functionspace::NodeColumns::mesh ) .def_property_readonly( "nodes", &functionspace::NodeColumns::nodes ) - .def_property_readonly( "valid", &functionspace::NodeColumns::valid ); + .def_property_readonly( "valid", &functionspace::NodeColumns::valid ) + .def( "__bool__", [](const functionspace::NodeColumns& self) { return self.valid(); } ); + py::class_( m_fs, "CellColumns" ) .def( py::init( []( Mesh const& m, int halo, int levels ) { return functionspace::CellColumns( m, util::Config()( "halo", halo )("levels",levels) ); @@ -489,10 +493,14 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "shape", py::overload_cast<>( &Field::shape, py::const_ ) ) .def_property_readonly( "size", &Field::size ) .def_property_readonly( "rank", &Field::rank ) + .def_property_readonly( "levels", &Field::levels ) .def_property_readonly( "datatype", []( Field& f ) { return atlasToPybind( f.datatype() ); } ) .def_property( "metadata", py::overload_cast<>( &Field::metadata, py::const_ ), py::overload_cast<>( &Field::metadata ) ) .def_property_readonly( "functionspace", py::overload_cast<>( &Field::functionspace, py::const_ ) ) + .def( "halo_exchange", []( Field& f) { f.haloExchange(); }) + .def_property_readonly( "halo_dirty", []( Field& f) { f.haloExchange(); }) + .def_property( "halo_dirty", &Field::dirty, &Field::set_dirty, py::return_value_policy::copy) .def_buffer( []( Field& f ) { auto strides = f.strides(); std::transform( strides.begin(), strides.end(), strides.begin(), @@ -532,11 +540,11 @@ PYBIND11_MODULE( _atlas4py, m ) { .def( "__exit__", []( output::Gmsh& gmsh, py::object exc_type, py::object exc_val, py::object exc_tb ) { gmsh.reset( nullptr ); } ) .def( - "write", []( output::Gmsh& gmsh, Mesh const& mesh ) { gmsh.write( mesh ); }, "mesh"_a ) + "write", []( output::Gmsh& gmsh, Mesh const& mesh ) { gmsh.write( mesh ); return gmsh; }, "mesh"_a ) .def( - "write", []( output::Gmsh& gmsh, Field const& field ) { gmsh.write( field ); }, "field"_a ) + "write", []( output::Gmsh& gmsh, Field const& field ) { gmsh.write( field ); return gmsh;}, "field"_a ) .def( - "write", []( output::Gmsh& gmsh, Field const& field, FunctionSpace const& fs ) { gmsh.write( field, fs ); }, + "write", []( output::Gmsh& gmsh, Field const& field, FunctionSpace const& fs ) { gmsh.write( field, fs ); return gmsh; }, "field"_a, "functionspace"_a ); From 6fb024a48dfa7569a45af9cef2229e3c95613793 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Tue, 7 Jun 2022 16:26:29 +0100 Subject: [PATCH 05/20] Add pyvista plotting, preliminary --- examples/example1.py | 12 ++++ examples/example2.py | 56 ++++++++++-------- src/atlas4py/__init__.py | 2 + src/atlas4py/pyvista/__init__.py | 97 ++++++++++++++++++++++++++++++++ 4 files changed, 145 insertions(+), 22 deletions(-) create mode 100644 src/atlas4py/pyvista/__init__.py diff --git a/examples/example1.py b/examples/example1.py index 4971c35..7a60366 100755 --- a/examples/example1.py +++ b/examples/example1.py @@ -21,5 +21,17 @@ gmsh.write(mesh) gmsh.write(field) +def plot_pyvista(mesh,field,title): + """ Plot with pyvista only mpi-serial for now""" + import pyvista + ds = atlas.pyvista.dataset_create(mesh, coordinates='xyz', field=field) + pyvista.set_plot_theme('paraview') + pl = pyvista.Plotter() + pl.add_mesh(ds, show_edges=True) + pl.add_title(title) + pl.show(interactive=True) + +# plot_pyvista(mesh,field,'title') + atlas.finalize() diff --git a/examples/example2.py b/examples/example2.py index a284e69..3f5d65a 100755 --- a/examples/example2.py +++ b/examples/example2.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 # Example program that initializes a spectral field with a -# spherical harmonic. Then the inverse transform is performed -# Then the result can be visualized with Gmsh: +# spherical harmonic. Then the inverse transform is performed. +# Then the result can be visualized with Gmsh and pyvista # gmsh example2.msh # ------------------------------------------------------------------------ @@ -25,7 +25,7 @@ # trans = atlas.Trans(grid, truncation) # instead of # trans = atlas.Trans(fs_gp, fs_sp) -# 2) Using atlas Field API, only level=1 is supported. +# 2) Using atlas Field API, only level=0 is supported (rank-1). # 3) The Healpix grid is identified as unstructured # because the grid wrongly gets cropped to a "west", and its # structure gets lost @@ -36,26 +36,37 @@ # ------------------------------------------------------------------------ # Functions +# ------------------------------------------------------------------------ def set_spherical_harmonic(field_sp, n, m): """ Set field_sp to a spherical harmonic """ - nflds = field_sp.shape[1] spectral_functionspace = atlas.functionspace.Spectral(field_sp.functionspace) if not spectral_functionspace: raise TypeError("'field_sp.functionspace' could not be cast to " + "'atlas.functionspace.Spectral'. It is of type " + "'atlas.functionspace."+field_sp.functionspace.type+"'.") - sp = atlas.make_view(field_sp) - sp[:,:] = 0. + + sp[:] = 0. def spherical_harmonic(re,im,_n,_m): - for jfld in range(nflds): - if n*(jfld+1) == _n and m*(jfld+1) == _m: - sp[re,jfld] = 1. - sp[im,jfld] = 1. + if _n == n and _m == m: + sp[re] = 1. + sp[im] = 1. spectral_functionspace.parallel_for( spherical_harmonic ) + # ------------------------------------------------------------------------ +def plot_pyvista(mesh,field,title): + """ Plot using pyviste, only works with mpi-serial for now """ + import pyvista + ds = atlas.pyvista.dataset_create(mesh, coordinates='xyz', field=field) + pyvista.set_plot_theme('paraview') + pl = pyvista.Plotter() + pl.add_mesh(ds, show_edges=False) + pl.add_title(title) + pl.show(interactive=True) + +# ------------------------------------------------------------------------ # ------------------------------------------------------------------------ # Start here @@ -63,41 +74,42 @@ def spherical_harmonic(re,im,_n,_m): atlas.initialize() # Required +init_total_wave_number=16 +init_zonal_wave_number=8 + grid = atlas.Grid("O128") -levels = 10 truncation = 255 atlas.Trans.backend("ectrans") partitioner = atlas.Partitioner("ectrans") # 'ectrans' partitioner needs to be used for 'ectrans' Trans backend only. - # It is essentially nearly identical to the equal_regions partitioner + # Partitions are nearly identical to the equal_regions partitioner if not atlas.GaussianGrid(grid) or not atlas.Trans.has_backend("ectrans"): # Fallback to "local" backend atlas.Trans.backend("local") - levels = 1 partitioner = atlas.Partitioner("equal_regions") # Create function spaces fs_sp = atlas.functionspace.Spectral(truncation) -fs_gp = atlas.functionspace.StructuredColumns(grid, partitioner) +fs_gp = atlas.functionspace.StructuredColumns(grid, partitioner, halo=1) # Create fields (already distributed) -field_sp = fs_sp.create_field(name="sp", levels=levels) -field_gp = fs_gp.create_field(name="gp", levels=levels) +field_sp = fs_sp.create_field(name="sp") +field_gp = fs_gp.create_field(name="gp") # Initial condition for spectral field -set_spherical_harmonic(field_sp, n=4, m=2) +set_spherical_harmonic(field_sp, n=init_total_wave_number, m=init_zonal_wave_number) # Spectral transform trans = atlas.Trans(grid, truncation) trans.invtrans(field_sp, field_gp) # Visualisation of gridpoint field -meshgenerator = atlas.MeshGenerator(type='structured') -mesh = meshgenerator.generate(grid) -gmsh = atlas.Gmsh("example2.msh", coordinates='xy') -gmsh.write(mesh) -gmsh.write(field_gp) +mesh = atlas.MeshGenerator(type='structured').generate(grid) +atlas.Gmsh("example2.msh", coordinates='xy').write(mesh).write(field_gp) + +# plot_pyvista(mesh, field_gp, title=grid.name + " - T" + str(truncation) + "; n="+str(init_total_wave_number)+" m="+str(init_zonal_wave_number)) +# --> only works with mpi-serial for now atlas.finalize() # ------------------------------------------------------------------------ diff --git a/src/atlas4py/__init__.py b/src/atlas4py/__init__.py index 5212c44..d36b5c1 100644 --- a/src/atlas4py/__init__.py +++ b/src/atlas4py/__init__.py @@ -7,6 +7,8 @@ from ._version import __version__ from ._atlas4py import * +from .pyvista import * + def make_view(field): import numpy as np return np.array(field, copy=False) diff --git a/src/atlas4py/pyvista/__init__.py b/src/atlas4py/pyvista/__init__.py new file mode 100644 index 0000000..7b180ca --- /dev/null +++ b/src/atlas4py/pyvista/__init__.py @@ -0,0 +1,97 @@ + +try: + import pyvista as pv + _atlas_pyvista_support = True +except ImportError: + _atlas_pyvista_support = False + +import numpy as np +import atlas4py as atlas + +def _atlas_pyvista_connectivity(conn): + num_neighbors = np.add.reduce(np.where(conn != -1, 1, 0), axis=1) + tmp = np.concatenate((num_neighbors[:, None], conn), axis=1) + return tmp[np.concatenate((np.full((conn.shape[0], 1), True, dtype=bool), conn != -1), axis=1)] + +def _atlas_connectivity_to_numpy(atlas_conn, *, out=None ): + if isinstance(atlas_conn, atlas.BlockConnectivity): + shape = (atlas_conn.rows, atlas_conn.cols) + out = np.zeros(shape, dtype=np.int64) if out is None else out + assert out.shape == shape + + for i in range(atlas_conn.rows): + for nb in range(atlas_conn.cols): + out[i, nb] = atlas_conn[i, nb] + + return out + + shape = (atlas_conn.rows, atlas_conn.maxcols) + out = np.zeros(shape, dtype=np.int64) if out is None else out + assert out.shape == shape + + for i in range(atlas_conn.rows): + cols = atlas_conn.cols(i) + for nb in range(cols): + out[i, nb] = atlas_conn[i, nb] + out[i, cols:]=-1 + + return out + +def _atlas_pyvista_cells(mesh): + np_conn = _atlas_connectivity_to_numpy(mesh.cells.node_connectivity) + return _atlas_pyvista_connectivity(np_conn) + +def _atlas_pyvista_points(mesh, coordinates='lonlat', **kwargs): + if coordinates=='lonlat': + return np.insert(mesh.nodes.lonlat, 2, 0., axis=1) + elif coordinates=='xy': + return np.insert(mesh.nodes.field('xy'), 2, 0., axis=1) + elif coordinates=='xyz': + deg2rad = np.pi/180. + xyrad = atlas.make_view(mesh.nodes.lonlat) * deg2rad + phi, theta = xyrad[:, 1], xyrad[:, 0] + return np.stack((np.cos(phi) * np.cos(theta), np.cos(phi) * np.sin(theta), np.sin(phi)), axis=1) + else: + raise IndexError("No coordinates "+coordinates+" supported") + +def _atlas_pyvista_mesh(mesh, coordinates): + import pyvista as pv + points = _atlas_pyvista_points( mesh, coordinates) + cells = _atlas_pyvista_cells(mesh) + return pv.PolyData(points, cells) + +def dataset_update(dataset,field,level=0,variable=0): + if atlas.functionspace.NodeColumns(field.functionspace): + field.halo_exchange() + if field.rank == 1: + dataset.point_data[field.name] = atlas.make_view(field)[:] + if field.rank == 2: + dataset.point_data[field.name] = atlas.make_view(field)[:,level] + if field.rank == 3: + dataset.point_data[field.name] = atlas.make_view(field)[:,level,variable] + else: + if not hasattr(dataset, 'atlas_mesh'): + raise Exception("no atlas_mesh stored. This is not a atlas-compatible pyvista dataset") + if not hasattr(dataset, 'nodes_fs'): + dataset.nodes_fs = atlas.functionspace.NodeColumns(dataset.atlas_mesh) + tmp_field = dataset.nodes_fs.create_field(dtype=field.datatype) + n_points = min(field.shape[0],dataset.number_of_points) + if field.rank == 1: + atlas.make_view(tmp_field)[:n_points] = atlas.make_view(field)[:n_points] + if field.rank == 2: + atlas.make_view(tmp_field)[:n_points] = atlas.make_view(field)[:n_points,level] + if field.rank == 3: + atlas.make_view(tmp_field)[:n_points] = atlas.make_view(field)[:n_points,level,variable] + tmp_field.halo_exchange() + dataset.point_data[field.name] = atlas.make_view(tmp_field) + +def dataset_create(mesh,coordinates='lonlat',field=None,level=0,variable=0): + if not _atlas_pyvista_support: + raise ModuleNotFoundError("No module named 'pyvista'") + + dataset = _atlas_pyvista_mesh(mesh, coordinates) + dataset.atlas_mesh = mesh + if field: + dataset_update(dataset,field=field,level=level,variable=variable) + return dataset + From b3669e5fd15e8abf0bb3e3625e4bde5c0646a2d1 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Mon, 27 Jun 2022 15:02:01 +0200 Subject: [PATCH 06/20] grid iterator --- src/atlas4py/_atlas4py.cpp | 43 +++++++++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index f6971af..e0a3e3b 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -18,6 +18,9 @@ #include "atlas/output/Gmsh.h" #include "atlas/library.h" #include "atlas/trans/Trans.h" +#include "atlas/interpolation.h" +#include "atlas/util/function/VortexRollup.h" +#include "atlas/util/function/SphericalHarmonic.h" #include "eckit/value/Value.h" #include "eckit/config/Configuration.h" @@ -149,6 +152,7 @@ array::DataType pybindToAtlas( py::dtype const& dtype ) { return { 0 }; } + } // namespace PYBIND11_MODULE( _atlas4py, m ) { @@ -211,9 +215,21 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "size", &Grid::size ) .def_property_readonly( "projection", &Grid::projection ) .def_property_readonly( "domain", &Grid::domain ) + .def( "lonlat", [](Grid const& self) { + return self.lonlat().begin(); + }) .def( "__repr__", []( Grid const& g ) { return "_atlas4py.Grid("_s + py::str( toPyObject( g.spec().get() ) ) + ")"_s; } ); + py::class_(m, "IteratorLonLat") + .def("__iter__", [](grid::IteratorLonLat& it) { return it; }) + .def("__next__", [](grid::IteratorLonLat& it) { + PointLonLat p; + if( !it.next(p) ) { + throw py::stop_iteration(); + } + return p;}); + py::class_( m, "Spacing" ) .def( "__len__", &grid::Spacing::size ) .def( "__getitem__", &grid::Spacing::operator[]) @@ -293,12 +309,17 @@ PYBIND11_MODULE( _atlas4py, m ) { .def( py::init( []( const std::string& type ) { return grid::Partitioner(type); } ) ) .def( py::init( [](){ return grid::Partitioner(); } ) ); + py::class_( m, "MatchingPartitioner" ) + .def( py::init( []( Mesh const& mesh, py::kwargs kwargs ) { return grid::MatchingPartitioner(mesh, to_config(kwargs)); } ) ) + .def( py::init( []( FunctionSpace const& functionspace, py::kwargs kwargs ) { return grid::MatchingPartitioner(functionspace, to_config(kwargs)); } ) ); + py::class_( m, "MeshGenerator" ) .def( py::init( []( py::kwargs kwargs ) { return MeshGenerator( to_config(kwargs)); } ) ) .def( py::init( []( util::Config const& config ) { return MeshGenerator( config ); } ) ) .def( py::init( []( const std::string& type ) { return MeshGenerator(type); } ) ) .def( py::init( [](){ return MeshGenerator(); } ) ) - .def( "generate", py::overload_cast( &MeshGenerator::generate, py::const_ ) ); + .def( "generate", [](MeshGenerator const& self, Grid const& grid){ return self.generate(grid); } ) + .def( "generate", [](MeshGenerator const& self, Grid const& grid, grid::Partitioner const& partitioner ){ return self.generate(grid, partitioner); } ); py::class_( m, "StructuredMeshGenerator" ) // TODO in FunctionSpace below we expose config options, not the whole config object @@ -308,6 +329,7 @@ PYBIND11_MODULE( _atlas4py, m ) { py::class_( m, "Mesh" ) .def( py::init( []( const Grid& grid ) { return Mesh(grid); } ) ) + .def( py::init( []( const Grid& grid, const grid::Partitioner& partitioner ) { return Mesh(grid,partitioner); } ) ) .def_property_readonly( "grid", &Mesh::grid ) .def_property_readonly( "projection", &Mesh::projection ) .def_property( "nodes", py::overload_cast<>( &Mesh::nodes, py::const_ ), py::overload_cast<>( &Mesh::nodes ) ) @@ -558,5 +580,24 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_static("backend", [] ( const std::string& backend ){ trans::Trans::backend(backend); } ) .def_static("has_backend", [] ( const std::string& backend ){ return trans::Trans::hasBackend(backend); } ); + py::class_( m, "Interpolation" ) + .def( py::init( [](const std::string& type, const FunctionSpace& source, const FunctionSpace& target, py::kwargs kwargs){ + auto config = to_config(kwargs); + config.set("type",type); + return Interpolation(config,source,target); + } ), "type"_a, "source"_a, "target"_a ) + .def( py::init( [](const std::string& type, const Grid& source, const Grid& target, py::kwargs kwargs){ + auto config = to_config(kwargs); + config.set("type",type); + return Interpolation(config,source,target); + } ), "type"_a, "source"_a, "target"_a ) + .def( "execute", []( Interpolation const& self, const Field& source, Field& target) { return self.execute(source,target);} ) + .def_property_readonly( "source", &Interpolation::source ) + .def_property_readonly( "target", &Interpolation::target ); + + + auto m_function = m.def_submodule( "function" ); + m_function.def("vortex_rollup", [](double lon, double lat, double t) { return util::function::vortex_rollup(lon,lat,t); } ); + m_function.def("spherical_harmonic", [](double lon, double lat,int n, int m ) { return util::function::spherical_harmonic(n,m,lon,lat); }, "lon"_a, "lat"_a, "n"_a, "m"_a ); } From e334bb1e3e8da655e4e4c1a985f9ec57da55b761 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Mon, 11 Sep 2023 10:33:40 +0200 Subject: [PATCH 07/20] Fixes --- examples/example2.py | 22 +++++++++++++++++++--- src/atlas4py/_atlas4py.cpp | 7 ++++--- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/examples/example2.py b/examples/example2.py index 3f5d65a..0b33725 100755 --- a/examples/example2.py +++ b/examples/example2.py @@ -52,6 +52,20 @@ def spherical_harmonic(re,im,_n,_m): sp[re] = 1. sp[im] = 1. + import numpy.random as random + random.seed(1) + rand1=random.rand(len(sp)) + random.seed(2) + rand2=random.rand(len(sp)) + def spectrum(re,im,_n,_m): + if _m==0: + sp[re] = pow(_n+1,-5/6)*(rand1[re]-0.5) + else: + sp[re] = pow(_n+1,-5/6)*(rand1[re]-0.5) + sp[im] = pow(_n+1,-5/6)*(rand2[im]-0.5) + + + #spectral_functionspace.parallel_for( spectrum ) spectral_functionspace.parallel_for( spherical_harmonic ) # ------------------------------------------------------------------------ @@ -77,8 +91,8 @@ def plot_pyvista(mesh,field,title): init_total_wave_number=16 init_zonal_wave_number=8 -grid = atlas.Grid("O128") -truncation = 255 +grid = atlas.Grid("F400") +truncation = 127 atlas.Trans.backend("ectrans") partitioner = atlas.Partitioner("ectrans") # 'ectrans' partitioner needs to be used for 'ectrans' Trans backend only. @@ -88,6 +102,8 @@ def plot_pyvista(mesh,field,title): # Fallback to "local" backend atlas.Trans.backend("local") partitioner = atlas.Partitioner("equal_regions") +atlas.Trans.backend("local") +partitioner = atlas.Partitioner("equal_regions") # Create function spaces fs_sp = atlas.functionspace.Spectral(truncation) @@ -106,7 +122,7 @@ def plot_pyvista(mesh,field,title): # Visualisation of gridpoint field mesh = atlas.MeshGenerator(type='structured').generate(grid) -atlas.Gmsh("example2.msh", coordinates='xy').write(mesh).write(field_gp) +atlas.Gmsh("example2.msh", coordinates='xyz').write(mesh).write(field_gp) # plot_pyvista(mesh, field_gp, title=grid.name + " - T" + str(truncation) + "; n="+str(init_total_wave_number)+" m="+str(init_zonal_wave_number)) # --> only works with mpi-serial for now diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index e0a3e3b..f4b55ea 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -203,13 +203,14 @@ PYBIND11_MODULE( _atlas4py, m ) { py::class_( m, "RectangularDomain" ) .def( py::init( []( std::tuple xInterval, std::tuple yInterval ) { auto [xFrom, xTo] = xInterval; - auto [yFrom, yTo] = xInterval; + auto [yFrom, yTo] = yInterval; return RectangularDomain( { xFrom, xTo }, { yFrom, yTo } ); } ), "x_interval"_a, "y_interval"_a ); py::class_( m, "Grid" ) .def( py::init() ) + .def( py::init( []( const std::string& name, const Domain& domain ) { return Grid(name,domain); } ) ) .def_property_readonly( "name", &Grid::name ) .def_property_readonly( "uid", &Grid::uid ) .def_property_readonly( "size", &Grid::size ) @@ -571,8 +572,8 @@ PYBIND11_MODULE( _atlas4py, m ) { py::class_( m, "Trans" ) - .def( py::init( [](const FunctionSpace& gp, const FunctionSpace& sp){ return trans::Trans(gp,sp);} ), "gp"_a, "sp"_a ) - .def( py::init( [](const Grid& grid, int truncation){ return trans::Trans(grid,truncation);} ), "grid"_a, "truncation"_a ) + .def( py::init( [](const FunctionSpace& gp, const FunctionSpace& sp, py::kwargs kwargs){ return trans::Trans(gp,sp,to_config(kwargs)); } ), "gp"_a, "sp"_a ) + .def( py::init( [](const Grid& grid, int truncation, py::kwargs kwargs){ return trans::Trans(grid,truncation,to_config(kwargs));} ), "grid"_a, "truncation"_a ) .def( "dirtrans", []( trans::Trans& trans, const Field& gpfield, Field& spfield) { trans.dirtrans(gpfield,spfield);} ) .def( "invtrans", []( trans::Trans& trans, const Field& spfield, Field& gpfield) { trans.invtrans(spfield,gpfield);} ) .def_property_readonly( "truncation", &trans::Trans::truncation ) From d72edf14bcff8c323f535110edb25945acf01552 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Mon, 11 Sep 2023 18:18:51 +0200 Subject: [PATCH 08/20] UnstructuredGrid --- src/atlas4py/_atlas4py.cpp | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index f4b55ea..bec4641 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -231,6 +231,34 @@ PYBIND11_MODULE( _atlas4py, m ) { } return p;}); + py::class_( m, "UnstructuredGrid" ) + .def( py::init( [](py::array_t _xy) { + py::buffer_info xy = _xy.request(); + auto xy_data = static_cast(xy.ptr); + auto points = new std::vector(xy.size/2); + auto& p = *points; + size_t j{0}; + for(size_t n=0; n _x, py::array_t _y) { + py::buffer_info x = _x.request(); + py::buffer_info y = _y.request(); + auto x_data = static_cast(x.ptr); + auto y_data = static_cast(x.ptr); + ATLAS_ASSERT(x.size == y.size); + auto points = new std::vector(x.size); + auto& p = *points; + for(size_t n=0; n( m, "Spacing" ) .def( "__len__", &grid::Spacing::size ) .def( "__getitem__", &grid::Spacing::operator[]) From 6672c48fdd001a2afc2bb7f05de851836927b489 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Tue, 19 Sep 2023 16:08:51 +0200 Subject: [PATCH 09/20] installation --- setup.py | 6 +++--- src/atlas4py/CMakeLists.txt | 1 + src/atlas4py/_atlas4py.cpp | 38 ++++++++++++++++++++++++++++--------- src/atlas4py/_version.py | 2 +- 4 files changed, 34 insertions(+), 13 deletions(-) diff --git a/setup.py b/setup.py index 2a95202..78a4656 100644 --- a/setup.py +++ b/setup.py @@ -13,8 +13,8 @@ cmake="3.14.0", ecbuild="3.6.3", eckit="1.17.1", - atlas="0.26.0", - pybind11="2.5.0", + atlas="develop", + pybind11="2.11.1", python="3.6", ) @@ -121,7 +121,7 @@ def build_extension(self, ext): subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=self.build_temp) -PACKAGE_VERSION = "0.29.0.dev14" +PACKAGE_VERSION = "0.34.0.dev14" # Meaning of the version scheme "{major}.{minor}.{patch}.dev{dev}": # - {major}.{minor}.{patch} => version of the atlas C++ library (hardcoded in 'setup.py') # - {dev} => version of the Python bindings as the commit number in 'master' diff --git a/src/atlas4py/CMakeLists.txt b/src/atlas4py/CMakeLists.txt index 3e5c71d..ee9347b 100644 --- a/src/atlas4py/CMakeLists.txt +++ b/src/atlas4py/CMakeLists.txt @@ -57,6 +57,7 @@ if(NOT atlas_FOUND) FetchContent_MakeAvailable(atlas) copy_target(atlas) + copy_target(atlas_io) if(_atlas4py_built_eckit) # copy atlas' eckit dependencies diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index bec4641..be2b3d8 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -234,27 +234,26 @@ PYBIND11_MODULE( _atlas4py, m ) { py::class_( m, "UnstructuredGrid" ) .def( py::init( [](py::array_t _xy) { py::buffer_info xy = _xy.request(); - auto xy_data = static_cast(xy.ptr); + auto xy_data = _xy.unchecked<2>(); auto points = new std::vector(xy.size/2); auto& p = *points; - size_t j{0}; for(size_t n=0; n _x, py::array_t _y) { py::buffer_info x = _x.request(); py::buffer_info y = _y.request(); - auto x_data = static_cast(x.ptr); - auto y_data = static_cast(x.ptr); + auto x_data = _x.unchecked<1>(); + auto y_data = _y.unchecked<1>(); ATLAS_ASSERT(x.size == y.size); auto points = new std::vector(x.size); auto& p = *points; for(size_t n=0; n( m_fs, "EdgeColumns" ) .def( py::init( []( Mesh const& m, int halo, int levels ) { return functionspace::EdgeColumns( m, util::Config()( "halo", halo )("levels", levels) ); @@ -513,6 +519,20 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "grid", &functionspace::StructuredColumns::grid ) .def_property_readonly( "valid", &functionspace::StructuredColumns::valid ); + py::class_( m_fs, "PointCloud" ) + .def( py::init( []( FunctionSpace fs ) { return functionspace::PointCloud{fs}; } ) ) + .def( py::init( []( Grid const& grid, py::kwargs kwargs) { + return functionspace::PointCloud( grid, to_config(kwargs) ); + } ), + "grid"_a ) + .def( py::init( []( Grid const& grid, atlas::grid::Partitioner const& partitioner, py::kwargs kwargs) { + return functionspace::PointCloud( grid, partitioner, to_config(kwargs) ); + } ), + "grid"_a, "partitioner"_a) + .def_property_readonly( "valid", &functionspace::PointCloud::valid ) + .def( "__bool__", [](const functionspace::PointCloud& self) { return self.valid(); } ); + + py::class_( m, "Metadata" ) .def_property_readonly( "keys", &util::Metadata::keys ) .def( "__setitem__", diff --git a/src/atlas4py/_version.py b/src/atlas4py/_version.py index f1cf221..c7a13ac 100644 --- a/src/atlas4py/_version.py +++ b/src/atlas4py/_version.py @@ -1 +1 @@ -__version__ = "0.26.0.dev14" +__version__ = "0.34.0.dev14" From 41f6105fe23ff9ad97de3a922e55b30b087832fe Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Thu, 21 Sep 2023 16:04:13 +0200 Subject: [PATCH 10/20] Add support for global fields and gather/scatter --- src/atlas4py/_atlas4py.cpp | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index be2b3d8..e2c9006 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -175,6 +175,7 @@ PYBIND11_MODULE( _atlas4py, m ) { "lon"_a, "lat"_a ) .def_property_readonly( "lon", py::overload_cast<>( &PointLonLat::lon, py::const_ ) ) .def_property_readonly( "lat", py::overload_cast<>( &PointLonLat::lat, py::const_ ) ) + .def( "__getitem__", &PointLonLat::operator() ) .def( "__repr__", []( PointLonLat const& p ) { return "_atlas4py.PointLonLat(lon=" + std::to_string( p.lon() ) + ", lat=" + std::to_string( p.lat() ) + ")"; @@ -186,6 +187,7 @@ PYBIND11_MODULE( _atlas4py, m ) { "x"_a, "y"_a ) .def_property_readonly( "x", py::overload_cast<>( &PointXY::x, py::const_ ) ) .def_property_readonly( "y", py::overload_cast<>( &PointXY::y, py::const_ ) ) + .def( "__getitem__", &PointLonLat::operator() ) .def( "__repr__", []( PointXY const& p ) { return "_atlas4py.PointXY(x=" + std::to_string( p.x() ) + ", y=" + std::to_string( p.y() ) + ")"; } ); @@ -460,13 +462,31 @@ PYBIND11_MODULE( _atlas4py, m ) { return fs.createField( config ); }, "dtype"_a = std::nullopt) + .def( + "create_field_global", + []( FunctionSpace const& fs, std::optional dtype, py::kwargs kwargs ) { + util::Config config = to_config(kwargs); + config.set("global",true); + if ( dtype ) + config.set( option::datatype( pybindToAtlas( py::dtype::from_args( *dtype ) ) )); + else + config.set( option::datatypeT() ); + return fs.createField( config ); + }, + "dtype"_a = std::nullopt) .def_property_readonly("lonlat", &FunctionSpace::lonlat ) .def_property_readonly("ghost", &FunctionSpace::ghost ) .def_property_readonly("remote_index", &FunctionSpace::remote_index ) .def_property_readonly("partition", &FunctionSpace::partition ) .def_property_readonly("global_index", &FunctionSpace::global_index ) .def_property_readonly("part", &FunctionSpace::part ) - .def_property_readonly("nb_parts", &FunctionSpace::nb_parts ); + .def_property_readonly("nb_parts", &FunctionSpace::nb_parts ) + .def("gather", [](FunctionSpace const& fs, Field const& local, Field& global) { + return fs.gather(local,global); + }) + .def("scatter", [](FunctionSpace const& fs, Field const& global, Field& local) { + return fs.scatter(global,local); + }); py::class_( m_fs, "EdgeColumns" ) .def( py::init( []( Mesh const& m, int halo, int levels ) { From df88101a981d7b58806ca878ac539615586c8cfd Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Thu, 21 Sep 2023 16:18:10 +0200 Subject: [PATCH 11/20] Add support for global fields and gather/scatter --- src/atlas4py/_atlas4py.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index e2c9006..ac70b9e 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -665,6 +665,14 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "target", &Interpolation::target ); + py::class_( m, "IndexKDTree" ) + .def( py::init([](){ return IndexKDTree(); })) + .def( "reserve", [](util::IndexKDTree& tree, idx_t size) { tree.reserve(size); }) + .def( "insert", [](util::IndexKDTree& tree, const util::IndexKDTree::Point& p, util::IndexKDTree::Payload )); + .def( "build", [](util::IndexKDTree& tree) { tree.build();} ) + ; + + auto m_function = m.def_submodule( "function" ); m_function.def("vortex_rollup", [](double lon, double lat, double t) { return util::function::vortex_rollup(lon,lat,t); } ); m_function.def("spherical_harmonic", [](double lon, double lat,int n, int m ) { return util::function::spherical_harmonic(n,m,lon,lat); }, "lon"_a, "lat"_a, "n"_a, "m"_a ); From 132892656a4bf24f349e616f2aedf6f7e553917c Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Fri, 22 Sep 2023 09:17:40 +0200 Subject: [PATCH 12/20] Python binding for IndexKDTree --- src/atlas4py/_atlas4py.cpp | 118 ++++++++++++++++++++++++++++++++----- 1 file changed, 104 insertions(+), 14 deletions(-) diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index ac70b9e..101f0a1 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -168,30 +168,46 @@ PYBIND11_MODULE( _atlas4py, m ) { .def("finalize", []() { atlas::finalize(); }) .def("finalise", []() { atlas::finalize(); }); - py::class_( m, "PointLonLat" ) + py::class_( m, "Point2" ) + .def( py::init( []( double x, double y ) { + return Point2( { x, y } ); + } ) ) + .def( "__getitem__", &Point2::operator() ) + .def( "__repr__", []( Point2 const& p ) { + return "_atlas4py.Point2(x=" + std::to_string( p.x() ) + ", y=" + std::to_string( p.y() ) + ")"; + } ); + + py::class_( m, "PointLonLat" ) .def( py::init( []( double lon, double lat ) { return PointLonLat( { lon, lat } ); } ), "lon"_a, "lat"_a ) .def_property_readonly( "lon", py::overload_cast<>( &PointLonLat::lon, py::const_ ) ) .def_property_readonly( "lat", py::overload_cast<>( &PointLonLat::lat, py::const_ ) ) - .def( "__getitem__", &PointLonLat::operator() ) .def( "__repr__", []( PointLonLat const& p ) { return "_atlas4py.PointLonLat(lon=" + std::to_string( p.lon() ) + ", lat=" + std::to_string( p.lat() ) + ")"; } ); - py::class_( m, "PointXY" ) - .def( py::init( []( double x, double y ) { - return PointXY( { x, y } ); - } ), - "x"_a, "y"_a ) - .def_property_readonly( "x", py::overload_cast<>( &PointXY::x, py::const_ ) ) - .def_property_readonly( "y", py::overload_cast<>( &PointXY::y, py::const_ ) ) - .def( "__getitem__", &PointLonLat::operator() ) + py::class_( m, "PointXY" ) .def( "__repr__", []( PointXY const& p ) { return "_atlas4py.PointXY(x=" + std::to_string( p.x() ) + ", y=" + std::to_string( p.y() ) + ")"; } ); + py::class_( m, "Point3" ) + .def( py::init( []( double x, double y, double z ) { + return Point3(x, y, z); + } ) ) + .def( "__getitem__", &Point3::operator() ) + .def( "__repr__", []( Point3 const& p ) { + return "_atlas4py.Point3(" + std::to_string( p(0) ) + ", " + std::to_string( p(1) ) + ", " + std::to_string( p(2) ) + ")"; + } ); + + py::class_( m, "PointXYZ" ) + .def( "__repr__", []( PointXYZ const& p ) { + return "_atlas4py.PointXYZ(x=" + std::to_string( p.x() ) + ", y=" + std::to_string( p.y() ) + ", z=" + std::to_string( p.z() ) + ")"; + } ); + + py::class_( m, "Projection" ).def( "__repr__", []( Projection const& p ) { return "_atlas4py.Projection("_s + py::str( toPyObject( p.spec().get() ) ) + ")"_s; } ); @@ -665,11 +681,85 @@ PYBIND11_MODULE( _atlas4py, m ) { .def_property_readonly( "target", &Interpolation::target ); + py::class_( m, "IndexKDTreeValue" ) + .def_property_readonly( "point", [](const util::IndexKDTree::Value& v) { return v.point(); } ) + .def_property_readonly( "payload", [](const util::IndexKDTree::Value& v) { return v.payload(); } ) + .def_property_readonly( "index", [](const util::IndexKDTree::Value& v) { return v.payload(); } ) + .def_property_readonly( "distance", [](const util::IndexKDTree::Value& v) { return v.distance(); } ) + ; + + py::class_( m, "IndexKDTreeValueList" ) + .def_property_readonly( "indices", [](const util::IndexKDTree::ValueList& v) { return v.payloads(); } ) + .def( "__getitem__", [](const util::IndexKDTree::ValueList& v, size_t i) { return v[i]; } ) + .def( "size", [](const util::IndexKDTree::ValueList& v) { return v.size(); } ) + ; + py::class_( m, "IndexKDTree" ) - .def( py::init([](){ return IndexKDTree(); })) - .def( "reserve", [](util::IndexKDTree& tree, idx_t size) { tree.reserve(size); }) - .def( "insert", [](util::IndexKDTree& tree, const util::IndexKDTree::Point& p, util::IndexKDTree::Payload )); - .def( "build", [](util::IndexKDTree& tree) { tree.build();} ) + .def( py::init([](){ return util::IndexKDTree(); })) + .def( "reserve", [](util::IndexKDTree& self, idx_t size) { self.reserve(size); }) + .def( "insert", [](util::IndexKDTree& self, const PointLonLat& point, util::IndexKDTree::Payload payload){ self.insert(point, payload); }) + .def( "insert", [](util::IndexKDTree& self, const PointXYZ& point, util::IndexKDTree::Payload payload){ self.insert(point, payload); }) + .def( "build", [](util::IndexKDTree& self) { self.build();} ) + .def( "build", [](util::IndexKDTree& self, py::array_t points) { + py::buffer_info info = points.request(); + auto size = info.shape[0]; + self.reserve(size); + if (info.ndim == 2 ) { + const PointLonLat* lonlat = reinterpret_cast(info.ptr); + for( size_t j=0; j(info.ptr); + for( size_t j=0; j Date: Fri, 22 Sep 2023 11:35:15 +0200 Subject: [PATCH 13/20] Remove erroneous field.halo_dirty property --- src/atlas4py/_atlas4py.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index 101f0a1..b0b08e3 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -606,7 +606,6 @@ PYBIND11_MODULE( _atlas4py, m ) { py::overload_cast<>( &Field::metadata ) ) .def_property_readonly( "functionspace", py::overload_cast<>( &Field::functionspace, py::const_ ) ) .def( "halo_exchange", []( Field& f) { f.haloExchange(); }) - .def_property_readonly( "halo_dirty", []( Field& f) { f.haloExchange(); }) .def_property( "halo_dirty", &Field::dirty, &Field::set_dirty, py::return_value_policy::copy) .def_buffer( []( Field& f ) { auto strides = f.strides(); From e51238b8b0ba1bebb4e084c3c8c2b370279955d1 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Fri, 22 Sep 2023 12:08:59 +0200 Subject: [PATCH 14/20] update ecbuild and eckit --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 78a4656..aaa2bea 100644 --- a/setup.py +++ b/setup.py @@ -11,8 +11,8 @@ VERSIONS = dict( cmake="3.14.0", - ecbuild="3.6.3", - eckit="1.17.1", + ecbuild="3.8.0", + eckit="1.24.0", atlas="develop", pybind11="2.11.1", python="3.6", From cbddf6acd045e6d44d393400ae99531b824baef6 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Fri, 22 Sep 2023 12:39:36 +0200 Subject: [PATCH 15/20] Example for unstructured pointcloud --- examples/example-unstructured.py | 340 +++++++++++++++++++++++++++++++ 1 file changed, 340 insertions(+) create mode 100755 examples/example-unstructured.py diff --git a/examples/example-unstructured.py b/examples/example-unstructured.py new file mode 100755 index 0000000..b5ca035 --- /dev/null +++ b/examples/example-unstructured.py @@ -0,0 +1,340 @@ +#!/usr/bin/env python3 +import atlas4py as atlas +import math +import numpy as np + +atlas.initialize() # Required, will also initialize MPI + +xy = np.array([ + [180,0], + [90,0], + [-90,0], + [0,90], + [0,-90], + [0,0], + [18,0], + [36,0], + [54,0], + [72,0], + [108,0], + [126,0], + [144,0], + [162,0], + [-162,0], + [-144,0], + [-126,0], + [-108,0], + [-72,0], + [-54,0], + [-36,0], + [-18,0], + [0,18], + [0,36], + [0,54], + [0,72], + [180,72], + [180,54], + [180,36], + [180,18], + [180,-18], + [180,-36], + [180,-54], + [180,-72], + [0,-72], + [0,-54], + [0,-36], + [0,-18], + [90,18], + [90,36], + [90,54], + [90,72], + [-90,72], + [-90,54], + [-90,36], + [-90,18], + [-90,-18], + [-90,-36], + [-90,-54], + [-90,-72], + [90,-72], + [90,-54], + [90,-36], + [90,-18], + [123.974,-58.6741], + [154.087,-16.9547], + [154.212,-58.8675], + [114.377,-41.9617], + [125.567,-23.5133], + [137.627,-40.8524], + [106.162,-24.5874], + [158.508,-38.55], + [137.826,-72.8109], + [142.103,-26.799], + [138.256,-13.8871], + [168.39,-24.3266], + [168.954,-12.0094], + [117.333,-12.35], + [102.254,-11.1537], + [120.307,59.7167], + [107.196,26.0167], + [144.768,28.3721], + [150.891,60.0343], + [164.566,25.5053], + [116.851,14.0295], + [124.84,28.3978], + [157.901,42.042], + [111.41,43.1056], + [134.333,44.6677], + [103.277,11.707], + [135.358,73.2119], + [135.349,14.2311], + [153.48,13.386], + [168.071,11.5344], + [-162.99,26.3775], + [-147.519,56.1313], + [-122.579,27.4824], + [-117.909,59.2376], + [-104.052,27.3616], + [-153.107,14.9717], + [-110.833,41.7436], + [-144.847,32.8534], + [-161.546,42.1031], + [-129.866,44.5201], + [-133.883,72.4163], + [-166.729,11.8907], + [-135.755,15.2529], + [-106.063,14.4869], + [-119.452,11.7037], + [-146.026,-58.6741], + [-115.913,-16.9547], + [-115.788,-58.8675], + [-155.623,-41.9617], + [-144.433,-23.5133], + [-132.373,-40.8524], + [-163.838,-24.5874], + [-111.492,-38.55], + [-132.174,-72.8109], + [-127.897,-26.799], + [-131.744,-13.8871], + [-101.61,-24.3266], + [-101.046,-12.0094], + [-152.667,-12.35], + [-167.746,-11.1537], + [-14.0127,-27.2963], + [-59.193,-57.0815], + [-56.465,-19.5751], + [-27.056,-59.3077], + [-57.124,-35.9752], + [-33.4636,-28.3914], + [-74.8037,-46.8602], + [-40.089,-45.1376], + [-74.8149,-28.3136], + [-21.3072,-42.2177], + [-44.0778,-72.6353], + [-19.6969,-12.8527], + [-40.1318,-12.1601], + [-72.691,-11.4129], + [-56.0261,58.6741], + [-25.9127,16.9547], + [-25.7876,58.8675], + [-65.6229,41.9617], + [-54.4335,23.5133], + [-42.373,40.8524], + [-73.838,24.5874], + [-21.4917,38.55], + [-42.1744,72.8109], + [-37.8974,26.799], + [-41.7437,13.8871], + [-11.6095,24.3266], + [-11.0459,12.0094], + [-62.667,12.35], + [-77.7456,11.1537], + [30.3071,59.7167], + [17.1956,26.0167], + [54.7676,28.3721], + [60.8915,60.0343], + [74.5657,25.5053], + [26.8506,14.0295], + [34.8398,28.3978], + [67.9014,42.042], + [21.41,43.1056], + [44.3335,44.6677], + [13.2772,11.707], + [45.3579,73.2119], + [45.3492,14.2311], + [63.4799,13.386], + [78.0714,11.5344], + [17.01,-26.3775], + [32.4806,-56.1313], + [57.4213,-27.4824], + [62.0912,-59.2376], + [75.9483,-27.3616], + [26.893,-14.9717], + [69.1672,-41.7436], + [35.1527,-32.8534], + [18.4543,-42.1031], + [50.1344,-44.5201], + [46.1172,-72.4163], + [13.2711,-11.8907], + [44.2448,-15.2529], + [73.9368,-14.4869], + [60.5478,-11.7037] +], dtype=np.float64) + + +# resolution = 2500 * 1000 # approximate ad hoc euclidian resolution in meters +# grid = atlas.UnstructuredGrid(xy) + +###  OR pass x,y individually +# grid = atlas.UnstructuredGrid(xy[:,0],xy[:,1]) + +### OR existing named grids + +resolution = 100 * 1000 # approximate euclidian resolution in meters +grid = atlas.Grid("H128") # HEALPix grid + +# grid = atlas.Grid("O320") # HEALPix grid +# resolution = 60 * 1000 # ad hoc + +# grid = atlas.Grid("O160") # HEALPix grid +# resolution = 120 * 1000 # ad hoc + + +### Create FunctionSpace + +functionspace = atlas.functionspace.PointCloud(grid, halo_radius=resolution*2) + +### Access parallelisation information, typically nb_parts == mpi_size, part == mpi_rank +# print("size", functionspace.size) +# print("nb_parts", functionspace.nb_parts) +# print("part", functionspace.part) + +### Access fields as numpy arrays, if needed +lonlat = atlas.make_view(functionspace.lonlat) # longitude latitude in degrees +ghost = atlas.make_view(functionspace.ghost) # ghost: 1 for ghost, 0 for owned +partition = atlas.make_view(functionspace.partition) # partition owning point (0-based) +remote_index = atlas.make_view(functionspace.remote_index) # local index on partition owning point (careful, 1-based when atlas_HAVE_FORTRAN) +global_index = atlas.make_view(functionspace.global_index) # global index across partitions (always 1-based) + +# internal = [ i for i in range(functionspace.size) if ghost[i] == 0 ] +# ghosts = [ i for i in range(functionspace.size) if ghost[i] == 1 ] + +### Create field and initialize +field = functionspace.create_field(name="myfield") +view = atlas.make_view(field) +for n in range(field.size): + if not ghost[n]: + lon = lonlat[n,0] * math.pi / 180. + lat = lonlat[n,1] * math.pi / 180. + view[n] = math.cos(4.*lat) * math.sin(8.*lon) + else: + view[n] = 0 + + +### Halo exchange +field.halo_dirty = True # if set to False, following halo_exchange will have no effect. Note it was already True upon create_field +field.halo_exchange() # comment to see halo with value 0, but validation below will fail +field.halo_exchange() # this will be ignored because field.halo_dirty was set to False during first call + + +### Validate halo exchange +validation_passed = True +for n in range(field.size): + if ghost[n]: + lon = lonlat[n,0] * math.pi / 180. + lat = lonlat[n,1] * math.pi / 180. + if view[n] != math.cos(4.*lat) * math.sin(8.*lon): + validation_passed = False + +if functionspace.part == 0: + if validation_passed: + print("validation of halo exchange passed") + else: + print("validation of halo exchange failed") + + +### Extra: search functionality +class Search: + def __init__(self, functionspace): + self.functionspace = functionspace + self.lonlat = atlas.make_view(self.functionspace.lonlat) + self.kdtree = atlas.IndexKDTree() + self.kdtree.build(lonlat) + + def nearest_indices_within_radius(self, i, radius): + # radius is in metres on Earth geometry + closest_indices = self.kdtree.closest_indices_within_radius(lon=self.lonlat[i,0], lat=self.lonlat[i,1], radius=radius) + return closest_indices[1:] # first point is "i" itself, so skip it + + +search = Search(functionspace) +radius = resolution +index = 0 +nearest = search.nearest_indices_within_radius(index,radius) +if functionspace.part == 0: + print("nearest global indices to local index",index," ( global index",global_index[index],"): ", [global_index[n] for n in nearest]) + +### Extra: scatter-plot global field, first doing a global gather on part 0 +def plot_global(field): + import matplotlib.pyplot as plt + import cartopy.crs as ccrs + + fs = field.functionspace + + ### Create global field and gather + field_global = fs.create_field_global(dtype=np.float64) + functionspace.gather(field, field_global) + + # plot on part 0 + if functionspace.part == 0: + x = np.zeros(grid.size) + y = np.zeros(grid.size) + for i,point in enumerate(grid.lonlat()): + x[i] = point.lon + y[i] = point.lat + z = atlas.make_view(field_global) + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree()) + + cntr = ax.scatter(x, y, c=z, s=1, cmap="RdBu_r", vmin=-1, vmax=1,transform=ccrs.PlateCarree()) + + ax.coastlines (resolution='110m', color='k') + ax.set_global() + ax.set_title('global field, %d points' % (grid.size) ) + ax.set_facecolor('gray') + fig.colorbar(cntr, ax=ax) + plt.subplots_adjust(hspace=0.5) + plt.show() + +### Extra: scatter-plot field of one partition +def plot_partition(field): + import matplotlib.pyplot as plt + import cartopy.crs as ccrs + + fs = field.functionspace + lonlat = atlas.make_view(fs.lonlat) + + x = lonlat[:,0] + y = lonlat[:,1] + z = atlas.make_view(field) + + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree()) + + cntr = ax.scatter(x, y, c=z, s=1, cmap="RdBu_r", vmin=-1, vmax=1,transform=ccrs.PlateCarree()) + + ax.coastlines (resolution='110m', color='k') + ax.set_global() + ax.set_title('part %d/%d , %d/%d points' % (fs.part, fs.nb_parts, field.size, grid.size) ) + ax.set_facecolor('gray') + fig.colorbar(cntr, ax=ax) + plt.subplots_adjust(hspace=0.5) + plt.show() + +# if functionspace.part == 0: # choose partition to plot here + # plot_partition(field) + +# plot_global(field) + +atlas.finalize() + From 93ea5a423cbbe32a2e3e7c9f65b08efb47357b7a Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Mon, 2 Oct 2023 12:30:40 +0200 Subject: [PATCH 16/20] atlas 0.35.0 released --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index aaa2bea..1111cb0 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ cmake="3.14.0", ecbuild="3.8.0", eckit="1.24.0", - atlas="develop", + atlas="0.35.0", pybind11="2.11.1", python="3.6", ) From 358e98b1ba5340812a5be1b992d48ceb378648a7 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Thu, 5 Oct 2023 10:59:58 +0200 Subject: [PATCH 17/20] Turn off gridtools_storage for now --- src/atlas4py/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/src/atlas4py/CMakeLists.txt b/src/atlas4py/CMakeLists.txt index ee9347b..a37fb76 100644 --- a/src/atlas4py/CMakeLists.txt +++ b/src/atlas4py/CMakeLists.txt @@ -54,6 +54,7 @@ if(NOT atlas_FOUND) GIT_REPOSITORY https://github.com/ecmwf/atlas.git GIT_TAG ${ATLAS4PY_ATLAS_VERSION} ) + set( ENABLE_GRIDTOOLS_STORAGE OFF CACHE BOOL "" FORCE ) FetchContent_MakeAvailable(atlas) copy_target(atlas) From 8a7396c48f25dbca396f4266012f6f6101850850 Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Thu, 5 Oct 2023 11:04:49 +0200 Subject: [PATCH 18/20] version 0.35.0.dev14 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 1111cb0..08146a0 100644 --- a/setup.py +++ b/setup.py @@ -121,7 +121,7 @@ def build_extension(self, ext): subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=self.build_temp) -PACKAGE_VERSION = "0.34.0.dev14" +PACKAGE_VERSION = "0.35.0.dev14" # Meaning of the version scheme "{major}.{minor}.{patch}.dev{dev}": # - {major}.{minor}.{patch} => version of the atlas C++ library (hardcoded in 'setup.py') # - {dev} => version of the Python bindings as the commit number in 'master' From 272a93c3347611a39b5811a7624acf4da91248ad Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Fri, 13 Oct 2023 17:26:54 +0200 Subject: [PATCH 19/20] Use unitsphere in example --- examples/example-unstructured.py | 15 ++++++++++----- src/atlas4py/_atlas4py.cpp | 1 + 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/examples/example-unstructured.py b/examples/example-unstructured.py index b5ca035..e5aa559 100755 --- a/examples/example-unstructured.py +++ b/examples/example-unstructured.py @@ -5,6 +5,10 @@ atlas.initialize() # Required, will also initialize MPI +# constants +earth_radius = 6371229. +km = 1000. + xy = np.array([ [180,0], [90,0], @@ -181,8 +185,8 @@ [60.5478,-11.7037] ], dtype=np.float64) - -# resolution = 2500 * 1000 # approximate ad hoc euclidian resolution in meters +# approximate ad hoc euclidian resolution in meters +# resolution = 2500 * km / earth_radius # grid = atlas.UnstructuredGrid(xy) ###  OR pass x,y individually @@ -190,7 +194,8 @@ ### OR existing named grids -resolution = 100 * 1000 # approximate euclidian resolution in meters +# approximate euclidian resolution in meters adimensionalized to unit sphere +resolution = 100 * km / earth_radius grid = atlas.Grid("H128") # HEALPix grid # grid = atlas.Grid("O320") # HEALPix grid @@ -202,7 +207,7 @@ ### Create FunctionSpace -functionspace = atlas.functionspace.PointCloud(grid, halo_radius=resolution*2) +functionspace = atlas.functionspace.PointCloud(grid, halo_radius=resolution*2, geometry="UnitSphere") ### Access parallelisation information, typically nb_parts == mpi_size, part == mpi_rank # print("size", functionspace.size) @@ -258,7 +263,7 @@ class Search: def __init__(self, functionspace): self.functionspace = functionspace self.lonlat = atlas.make_view(self.functionspace.lonlat) - self.kdtree = atlas.IndexKDTree() + self.kdtree = atlas.IndexKDTree(geometry="UnitSphere") self.kdtree.build(lonlat) def nearest_indices_within_radius(self, i, radius): diff --git a/src/atlas4py/_atlas4py.cpp b/src/atlas4py/_atlas4py.cpp index b0b08e3..9dddfb7 100644 --- a/src/atlas4py/_atlas4py.cpp +++ b/src/atlas4py/_atlas4py.cpp @@ -695,6 +695,7 @@ PYBIND11_MODULE( _atlas4py, m ) { py::class_( m, "IndexKDTree" ) .def( py::init([](){ return util::IndexKDTree(); })) + .def( py::init([](const std::string& geometry){ return util::IndexKDTree(util::Config("geometry",geometry)); }), "geometry"_a) .def( "reserve", [](util::IndexKDTree& self, idx_t size) { self.reserve(size); }) .def( "insert", [](util::IndexKDTree& self, const PointLonLat& point, util::IndexKDTree::Payload payload){ self.insert(point, payload); }) .def( "insert", [](util::IndexKDTree& self, const PointXYZ& point, util::IndexKDTree::Payload payload){ self.insert(point, payload); }) From c4048e553b12bfac2b520aae08879c36c58a5d5e Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Tue, 24 Oct 2023 22:14:34 +0200 Subject: [PATCH 20/20] update to version 0.35.1 needed for kdtree search with unit-sphere --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 08146a0..aec8952 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ cmake="3.14.0", ecbuild="3.8.0", eckit="1.24.0", - atlas="0.35.0", + atlas="0.35.1", pybind11="2.11.1", python="3.6", ) @@ -121,7 +121,7 @@ def build_extension(self, ext): subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=self.build_temp) -PACKAGE_VERSION = "0.35.0.dev14" +PACKAGE_VERSION = "0.35.1.dev14" # Meaning of the version scheme "{major}.{minor}.{patch}.dev{dev}": # - {major}.{minor}.{patch} => version of the atlas C++ library (hardcoded in 'setup.py') # - {dev} => version of the Python bindings as the commit number in 'master'