From 65e5c051a1c8441039c12a81e6ab2378a3edcabe Mon Sep 17 00:00:00 2001 From: Mike Kryjak Date: Wed, 29 Apr 2026 15:56:38 +0100 Subject: [PATCH] clang-format all VANTAGE coupling files --- bout-particle-push.cxx | 161 +-- bout-particle-push.hxx | 35 +- .../density_history.py | 4 +- ...nd_global_vertices_save_to_hypnotoad_nc.py | 15 +- .../load_and_plot_dm.py | 27 +- .../particle_animator.py | 79 +- .../petsc_dmplex_from_hypnotoad.py | 92 +- .../plot_corners_functions.py | 992 +++++++++++------- .../dmplex-vertex-coordinates/runtest | 39 +- tests/integrated/particle-pusher/runtest | 273 +++-- 10 files changed, 1033 insertions(+), 684 deletions(-) diff --git a/bout-particle-push.cxx b/bout-particle-push.cxx index 3a5c4e95a..86059497d 100644 --- a/bout-particle-push.cxx +++ b/bout-particle-push.cxx @@ -1,10 +1,10 @@ +#include "bout-particle-push.hxx" #include "bout/bout.hxx" #include "bout/bout_types.hxx" #include "bout/field2d.hxx" #include "bout/output.hxx" #include "bout/petsclib.hxx" #include -#include "bout-particle-push.hxx" #include #include #include @@ -23,8 +23,8 @@ #include #include // for reactions integration -#include #include "include/vantage_dmplex.hxx" +#include #ifndef NESO_PARTICLES_PETSC static_assert(false, "NESO-Particles was installed without PETSc support."); @@ -73,7 +73,6 @@ void calculate_neutral_density_in_place( // extrapolate -> Neumann } - double calculate_total_mass(Field2D& density, std::shared_ptr& neso_mesh) { double local_mass = 0.0; @@ -192,18 +191,25 @@ void check_cell_volumes(std::shared_ptr& neso_m } } -REAL cell_length(std::vector> cell_vertices, INT iv1, INT iv2, INT iv3, INT iv4){ - const REAL Rlength2 = std::pow(0.5*(cell_vertices.at(iv1).at(0) + cell_vertices.at(iv2).at(0) - - cell_vertices.at(iv3).at(0) - cell_vertices.at(iv4).at(0)),2.0); - const REAL Zlength2 = std::pow(0.5*(cell_vertices.at(iv1).at(1) + cell_vertices.at(iv2).at(1) - - cell_vertices.at(iv3).at(1) - cell_vertices.at(iv4).at(1)),2.0); - REAL length = std::pow(Zlength2+Rlength2,0.5); +REAL cell_length(std::vector> cell_vertices, INT iv1, INT iv2, INT iv3, + INT iv4) { + const REAL Rlength2 = + std::pow(0.5 + * (cell_vertices.at(iv1).at(0) + cell_vertices.at(iv2).at(0) + - cell_vertices.at(iv3).at(0) - cell_vertices.at(iv4).at(0)), + 2.0); + const REAL Zlength2 = + std::pow(0.5 + * (cell_vertices.at(iv1).at(1) + cell_vertices.at(iv2).at(1) + - cell_vertices.at(iv3).at(1) - cell_vertices.at(iv4).at(1)), + 2.0); + REAL length = std::pow(Zlength2 + Rlength2, 0.5); return length; } - void check_cell_centres(std::shared_ptr& neso_mesh, - Mesh*& bout_mesh, BoutReal absolute_tolerance, BoutReal relative_tolerance) { + Mesh*& bout_mesh, BoutReal absolute_tolerance, + BoutReal relative_tolerance) { Coordinates* coord = bout_mesh->getCoordinates(); // get (R,Z) of cell centres in Hypnotoad grid Field2D Rxy; @@ -221,7 +227,7 @@ void check_cell_centres(std::shared_ptr& neso_m REAL neso_Rxy = 0.0; REAL neso_Zxy = 0.0; - for (PetscInt iv = 0; iv < 4; iv++){ + for (PetscInt iv = 0; iv < 4; iv++) { neso_Rxy += cell_vertices.at(iv).at(0); neso_Zxy += cell_vertices.at(iv).at(1); } @@ -233,23 +239,25 @@ void check_cell_centres(std::shared_ptr& neso_m const REAL min_cell_length = std::min(cell_length_a, cell_length_b); // we compare the difference in cell centres to the absolute tolerance and // the relative tolerance formed by comparing to the smallest length across the cell - const REAL tolerance = absolute_tolerance + min_cell_length*relative_tolerance; - const bool centres_match = (abs(neso_Rxy - bout_Rxy) < tolerance) && - (abs(neso_Zxy - bout_Zxy) < tolerance); + const REAL tolerance = absolute_tolerance + min_cell_length * relative_tolerance; + const bool centres_match = (abs(neso_Rxy - bout_Rxy) < tolerance) + && (abs(neso_Zxy - bout_Zxy) < tolerance); // exit if we fail to find a match - NESOASSERT(centres_match, - fmt::format("Hypnotoad/BOUT++ cell centre (R, Z) ({}, {}) does not match NESO-Particles mesh " - "cell centre ({}, {}) for ix = {} iy = {} \n" - "The cell height and width are {} {} \n" - "The displacements in R and Z are {} {} \n" - "Ignore this message by " - "setting [neso_particles] test_cell_centres = false\n Relax the " - "tolerance used in this check by increasing\n" - "[neso_particles] cell_centre_absolute_tolerance = {}\n" - "[neso_particles] cell_centre_relative_tolerance = {}", - bout_Rxy, bout_Zxy, neso_Rxy, neso_Zxy, ix, iy, - cell_length_a, cell_length_b, abs(neso_Rxy - bout_Rxy), abs(neso_Zxy - bout_Zxy), - absolute_tolerance, relative_tolerance)); + NESOASSERT( + centres_match, + fmt::format("Hypnotoad/BOUT++ cell centre (R, Z) ({}, {}) does not match " + "NESO-Particles mesh " + "cell centre ({}, {}) for ix = {} iy = {} \n" + "The cell height and width are {} {} \n" + "The displacements in R and Z are {} {} \n" + "Ignore this message by " + "setting [neso_particles] test_cell_centres = false\n Relax the " + "tolerance used in this check by increasing\n" + "[neso_particles] cell_centre_absolute_tolerance = {}\n" + "[neso_particles] cell_centre_relative_tolerance = {}", + bout_Rxy, bout_Zxy, neso_Rxy, neso_Zxy, ix, iy, cell_length_a, + cell_length_b, abs(neso_Rxy - bout_Rxy), abs(neso_Zxy - bout_Zxy), + absolute_tolerance, relative_tolerance)); ixy++; } } @@ -270,8 +278,7 @@ void check_mass_conservation(double total_mass_final, double total_mass_initial, } VantageSourceManager::VantageSourceManager( - std::shared_ptr& neso_mesh, - Mesh* bout_mesh, + std::shared_ptr& neso_mesh, Mesh* bout_mesh, const std::map& norms) { // Store DMPlex and BOUT++ mesh @@ -354,12 +361,12 @@ int main(int argc, char** argv) { std::string dmplex_h5_filename = Options::root()["mesh"]["dmplex_h5_filename"].withDefault( "hypnotoad_dmplex_mesh_output.h5"); - DM dm = create_dmplex_from_Bout_mesh(bout_mesh, sycl_target, make_output_path(dmplex_h5_filename)); + DM dm = create_dmplex_from_Bout_mesh(bout_mesh, sycl_target, + make_output_path(dmplex_h5_filename)); output << "Begin particle push \n"; // get data from BOUT.inp to assign particle weights as a fn of x,y auto& opt = Options::root(); - /* * * @@ -402,17 +409,21 @@ int main(int argc, char** argv) { const int npart_per_cell = Options::root()["VANTAGE_reactions"]["npart_per_cell"].withDefault(1); - // Plasma parameters - const BoutReal background_ion_temperature = opt["VANTAGE_reactions"]["background_ion_temperature"].withDefault(1.0); - const BoutReal background_ion_density = opt["VANTAGE_reactions"]["background_ion_density"].withDefault(1.0); - const BoutReal background_ion_Vx = opt["VANTAGE_reactions"]["background_ion_Vx"].withDefault(0.0); - const BoutReal background_ion_Vy = opt["VANTAGE_reactions"]["background_ion_Vy"].withDefault(0.0); + const BoutReal background_ion_temperature = + opt["VANTAGE_reactions"]["background_ion_temperature"].withDefault(1.0); + const BoutReal background_ion_density = + opt["VANTAGE_reactions"]["background_ion_density"].withDefault(1.0); + const BoutReal background_ion_Vx = + opt["VANTAGE_reactions"]["background_ion_Vx"].withDefault(0.0); + const BoutReal background_ion_Vy = + opt["VANTAGE_reactions"]["background_ion_Vy"].withDefault(0.0); const std::vector V_background = {background_ion_Vx, background_ion_Vy}; // Reaction settings const REAL iz_rate = Options::root()["VANTAGE_reactions"]["iz_rate"].withDefault(1.0); - const REAL rec_rate = Options::root()["VANTAGE_reactions"]["rec_rate"].withDefault(1.0); + const REAL rec_rate = + Options::root()["VANTAGE_reactions"]["rec_rate"].withDefault(1.0); const BoutReal rec_markers_per_cell = Options::root()["VANTAGE_reactions"]["rec_markers_per_cell"].withDefault(1000); @@ -424,7 +435,6 @@ int main(int argc, char** argv) { .doc("Number of RNG samples to prepare per-particle") .withDefault(40); - BoutReal sim_time = 0.0; Field2D ion_density = Field2D(background_ion_density, bout_mesh); Field2D neutral_density = Field2D(0.0, bout_mesh); @@ -443,13 +453,15 @@ int main(int argc, char** argv) { if (Options::root()["neso_particles"]["test_cell_volumes"].withDefault(true)) { check_cell_volumes(neso_mesh, bout_mesh); } - if (Options::root()["neso_particles"]["test_cell_centres"].withDefault(true)){ - check_cell_centres(neso_mesh,bout_mesh, - Options::root()["neso_particles"]["cell_centre_absolute_tolerance"].withDefault(1.0e-12), - Options::root()["neso_particles"]["cell_centre_relative_tolerance"].withDefault(0.0)); + if (Options::root()["neso_particles"]["test_cell_centres"].withDefault(true)) { + check_cell_centres( + neso_mesh, bout_mesh, + Options::root()["neso_particles"]["cell_centre_absolute_tolerance"].withDefault( + 1.0e-12), + Options::root()["neso_particles"]["cell_centre_relative_tolerance"].withDefault( + 0.0)); } - // create a Reactions particle spec auto particle_spec_builder = ParticleSpecBuilder(ndim); auto electron_species = Species("ELECTRON"); @@ -532,7 +544,6 @@ int main(int argc, char** argv) { // Ionisation reaction // ------------------------------------------------------------------------------ - auto iz_rate_data = FixedRateData(iz_rate); main_species.set_id(0); auto ionisation_reaction = ElectronImpactIonisation( @@ -547,7 +558,6 @@ int main(int argc, char** argv) { // Options and constants - // Make new particle group just for the markers auto marker_group = std::make_shared(domain, particle_spec, sycl_target); @@ -583,11 +593,9 @@ int main(int argc, char** argv) { // Calculate marker weights - // Add particle property: number of particles in the local cell // From demo app: "distribute_n_part_cell" - for (int ic = 0; ic < marker_group->domain->mesh->get_cell_count(); ic++) - { + for (int ic = 0; ic < marker_group->domain->mesh->get_cell_count(); ic++) { int n_part_cell = marker_group->get_npart_cell(ic); particle_loop( "Update N_CELL prop", marker_group, @@ -606,39 +614,37 @@ int main(int argc, char** argv) { particle_loop( "Update weight of ions", marker_group, [=](auto n_cell_prop, auto ion_dens_prop, auto weight_prop) { - auto updated_weight = (ion_dens_prop.at(0) * Nnorm * V_cell) / - (N_w * n_cell_prop.at(0)); + auto updated_weight = + (ion_dens_prop.at(0) * Nnorm * V_cell) / (N_w * n_cell_prop.at(0)); weight_prop.at(0) = updated_weight; }, - Access::read(Sym("N_CELL")), - Access::read(Sym("FLUID_DENSITY")), + Access::read(Sym("N_CELL")), Access::read(Sym("FLUID_DENSITY")), Access::write(Sym("WEIGHT"))) ->execute(ic); } // Define marker species and reaction rates - auto recomb_species = Species("ION", 1.0, 0.0, -1); //TODO: better as marker_species + auto recomb_species = Species("ION", 1.0, 0.0, -1); // TODO: better as marker_species auto recomb_data = FixedRateData(rec_rate); - auto recomb_energy_data = FixedRateData(rec_rate); //TODO: make this separate + auto recomb_energy_data = FixedRateData(rec_rate); // TODO: make this separate // This sampler will calculate marker momentum from fluid plasma conditions // TODO: Do I need a separate rng kernel? auto constant_rate_cross_section = ConstantRateCrossSection(1.0); auto recomb_data_calc_sampler = - FilteredMaxwellianSampler<2, decltype(constant_rate_cross_section)>( - 1 / (recomb_species.get_mass() * Tnorm), - constant_rate_cross_section, rng_kernel); + FilteredMaxwellianSampler<2, decltype(constant_rate_cross_section)>( + 1 / (recomb_species.get_mass() * Tnorm), constant_rate_cross_section, + rng_kernel); // Container for objects allowing calculation of parameters within // the recombination kernel: sampled velocity and the radiation // energy loss source. Must be in this order. auto recomb_data_calc_obj = - DataCalculator( - recomb_energy_data, recomb_data_calc_sampler); + DataCalculator( + recomb_energy_data, recomb_data_calc_sampler); - BoutReal normalised_potential_energy = 1.0; //TODO: units + BoutReal normalised_potential_energy = 1.0; // TODO: units auto recomb_reaction_kernel = RecombReactionKernels<2>( recomb_species, electron_species, normalised_potential_energy); @@ -647,13 +653,11 @@ int main(int argc, char** argv) { std::array recomb_out_states = {out_state}; // Create reaction object - auto recomb_reaction = LinearReactionBase<1, decltype(recomb_data), - decltype(recomb_reaction_kernel), - decltype(recomb_data_calc_obj)>( - sycl_target, recomb_species.get_id(), recomb_out_states, - recomb_data, recomb_reaction_kernel, recomb_data_calc_obj); - - + auto recomb_reaction = + LinearReactionBase<1, decltype(recomb_data), decltype(recomb_reaction_kernel), + decltype(recomb_data_calc_obj)>( + sycl_target, recomb_species.get_id(), recomb_out_states, recomb_data, + recomb_reaction_kernel, recomb_data_calc_obj); // Wrappers & controllers // ------------------------------------------------------------------------------ @@ -685,8 +689,9 @@ int main(int argc, char** argv) { auto accumulator_real_transform_wrapper = std::make_shared( std::dynamic_pointer_cast(accumulator_transform_iz)); - auto ion_source_density_zeroer = make_transformation_strategy>( - std::vector{"ION_SOURCE_DENSITY"}); + auto ion_source_density_zeroer = + make_transformation_strategy>( + std::vector{"ION_SOURCE_DENSITY"}); std::vector> child_transforms = std::vector{merge_wrapper, remove_wrapper}; @@ -711,8 +716,8 @@ int main(int argc, char** argv) { std::vector> parent_transforms_rec = std::vector{accumulator_real_transform_wrapper, merge_wrapper, remove_wrapper}; - auto recombination_controller = ReactionController( - parent_transforms_rec, child_transforms); + auto recombination_controller = + ReactionController(parent_transforms_rec, child_transforms); recombination_controller.add_reaction( std::make_shared(recomb_reaction)); @@ -794,8 +799,8 @@ int main(int argc, char** argv) { } }; // uncomment to write a trajectory - H5Part h5part(make_output_path("particle_trajectories.h5part"), - A_particle_group, Sym("POSITION"), Sym("VELOCITY")); + H5Part h5part(make_output_path("particle_trajectories.h5part"), A_particle_group, + Sym("POSITION"), Sym("VELOCITY")); // allocate buffer vector for scalar projection/evaluation of NESO-Particles // properties @@ -805,7 +810,8 @@ int main(int argc, char** argv) { neso_mesh, h_project1); // Calculate neutral density and sources for initial condition - calculate_neutral_density_in_place(neutral_density, dg0, A_particle_group, h_project1); + calculate_neutral_density_in_place(neutral_density, dg0, A_particle_group, + h_project1); source_manager.update_all_sources(dt); // diagnose the initial condition @@ -831,7 +837,8 @@ int main(int argc, char** argv) { // uncomment to write a trajectory h5part.write(); - calculate_neutral_density_in_place(neutral_density, dg0, A_particle_group, h_project1); + calculate_neutral_density_in_place(neutral_density, dg0, A_particle_group, + h_project1); source_manager.update_all_sources(dt); Field2D Siz = source_manager.get_data("Siz"); Field2D Srec = source_manager.get_data("Srec"); diff --git a/bout-particle-push.hxx b/bout-particle-push.hxx index ad8d4cd2b..9628e19e7 100644 --- a/bout-particle-push.hxx +++ b/bout-particle-push.hxx @@ -9,7 +9,8 @@ using namespace VANTAGE::Reactions; /// @brief Data struct to hold information about a reaction source. /// @param reaction_name Name of the reaction, e.g. "ionistaion" -/// @param source_name Name of the source, e.g. Siz (ion density source due to ionisation). +/// @param source_name Name of the source, e.g. Siz (ion density source due to +/// ionisation). /// @param accumulator CellwiseAccumulator to use to accumulate the source term for this /// reaction. /// @param particle_group ParticleGroup to which this source applies. @@ -31,8 +32,7 @@ struct VantageSource { class VantageSourceManager { public: VantageSourceManager(std::shared_ptr& neso_mesh, - Mesh* bout_mesh, - const std::map& norms); + Mesh* bout_mesh, const std::map& norms); Mesh* bout_mesh; @@ -70,8 +70,6 @@ private: * */ - - /** * @brief Function to calculate particle positions and velocities from a Maxwellian. * @@ -87,10 +85,12 @@ private: */ template -inline ParticleSet uniform_cellwise_maxwellian( - SYCLTargetSharedPtr sycl_target, std::shared_ptr mesh, - const ParticleSpec &particle_spec, const INT &npart_per_cell, - const REAL &weight, const REAL &std_dev, const INT &species_id) { +inline ParticleSet +uniform_cellwise_maxwellian(SYCLTargetSharedPtr sycl_target, + std::shared_ptr mesh, + const ParticleSpec& particle_spec, const INT& npart_per_cell, + const REAL& weight, const REAL& std_dev, + const INT& species_id) { const int rank = sycl_target->comm_pair.rank_parent; const int size = sycl_target->comm_pair.size_parent; @@ -101,12 +101,11 @@ inline ParticleSet uniform_cellwise_maxwellian( std::vector> positions; std::vector cell_ids; PetscInterface::uniform_within_dmplex_cells(mesh, npart_per_cell, positions, cell_ids, - &rng_pos); + &rng_pos); const int N = cell_ids.size(); - auto velocities = - NESO::Particles::normal_distribution(N, ndim, 0.0, std_dev, rng_vel); + auto velocities = NESO::Particles::normal_distribution(N, ndim, 0.0, std_dev, rng_vel); ParticleSet maxwellian(N, particle_spec); @@ -129,14 +128,13 @@ inline ParticleSet uniform_cellwise_maxwellian( * in recombination and charge exchange. */ -inline auto get_uniform_rng_kernel(SYCLTargetSharedPtr sycl_target, - std::size_t n_samples, +inline auto get_uniform_rng_kernel(SYCLTargetSharedPtr sycl_target, std::size_t n_samples, std::uint64_t root_seed = 141351) { const int rank = sycl_target->comm_pair.rank_parent; - std::uint64_t seed = NESO::RNGToolkit::create_seeds( - sycl_target->comm_pair.size_parent, rank, root_seed); + std::uint64_t seed = + NESO::RNGToolkit::create_seeds(sycl_target->comm_pair.size_parent, rank, root_seed); auto rng_normal = NESO::RNGToolkit::create_rng( NESO::RNGToolkit::Distribution::Uniform{ @@ -146,12 +144,11 @@ inline auto get_uniform_rng_kernel(SYCLTargetSharedPtr sycl_target, // Create an interface between NESO-RNG-Toolkit and NESO-Particles KernelRNG auto rng_interface = make_rng_generation_function( - [=](REAL *d_ptr, const std::size_t num_samples) -> int { + [=](REAL* d_ptr, const std::size_t num_samples) -> int { return rng_normal->get_samples(d_ptr, num_samples); }); - auto rng_kernel = - host_atomic_block_kernel_rng(rng_interface, n_samples); + auto rng_kernel = host_atomic_block_kernel_rng(rng_interface, n_samples); return rng_kernel; } \ No newline at end of file diff --git a/scripts/dmplex_tools_for_particle_pusher/density_history.py b/scripts/dmplex_tools_for_particle_pusher/density_history.py index 54d94b607..27ddfd8ef 100644 --- a/scripts/dmplex_tools_for_particle_pusher/density_history.py +++ b/scripts/dmplex_tools_for_particle_pusher/density_history.py @@ -7,7 +7,6 @@ def plot(case_path): - ds = xbout.load.open_boutdataset( datapath=case_path / "BOUT.dmp.*.nc", inputfilepath=case_path / "BOUT.inp", @@ -27,7 +26,7 @@ def plot(case_path): Ni_avg.plot(ax=ax, marker="o") ax.set_ylabel("Normalised density") ax.set_title("Mean ion density") - + ax = axes[1] Nn_avg.plot(ax=ax, marker="o") ax.set_ylabel("Normalised density") @@ -45,7 +44,6 @@ def plot(case_path): ax.set_ylabel("-") ax.set_title("Density volume integrals") - for ax in axes: ax.set_xlabel("time") diff --git a/scripts/dmplex_tools_for_particle_pusher/find_global_vertices_save_to_hypnotoad_nc.py b/scripts/dmplex_tools_for_particle_pusher/find_global_vertices_save_to_hypnotoad_nc.py index 9d9e33730..ae70b019f 100644 --- a/scripts/dmplex_tools_for_particle_pusher/find_global_vertices_save_to_hypnotoad_nc.py +++ b/scripts/dmplex_tools_for_particle_pusher/find_global_vertices_save_to_hypnotoad_nc.py @@ -1,9 +1,18 @@ from plot_corners_functions import plot_corners_get_dmplex_data import argparse -parser = argparse.ArgumentParser(description="Process a Hypnotoad mesh file to produce a modified mesh file and a file containing the global list of vertices of the mesh. Produce plots of cell vertices.") -parser.add_argument("hypnotoad_nc_file_path", type=str, help="The path to the Hypnotoad netCDF file representing the mesh.") + +parser = argparse.ArgumentParser( + description="Process a Hypnotoad mesh file to produce a modified mesh file and a file containing the global list of vertices of the mesh. Produce plots of cell vertices." +) +parser.add_argument( + "hypnotoad_nc_file_path", + type=str, + help="The path to the Hypnotoad netCDF file representing the mesh.", +) args = parser.parse_args() file_path = args.hypnotoad_nc_file_path print(f"Processing Hypnotoad mesh from {args.hypnotoad_nc_file_path}") -Nx,Ny,cell_vertices,vertex_list,boundary_vertex_info = plot_corners_get_dmplex_data(file_path,interactive_plot=True,print_cells_to_screen_output=False) +Nx, Ny, cell_vertices, vertex_list, boundary_vertex_info = plot_corners_get_dmplex_data( + file_path, interactive_plot=True, print_cells_to_screen_output=False +) diff --git a/scripts/dmplex_tools_for_particle_pusher/load_and_plot_dm.py b/scripts/dmplex_tools_for_particle_pusher/load_and_plot_dm.py index dc388bba4..bc3a99f9d 100644 --- a/scripts/dmplex_tools_for_particle_pusher/load_and_plot_dm.py +++ b/scripts/dmplex_tools_for_particle_pusher/load_and_plot_dm.py @@ -1,14 +1,19 @@ from petsc4py import PETSc import argparse + parser = argparse.ArgumentParser(description="Plot the edges of a DMPlex mesh.") -parser.add_argument("dmplex_h5_file_path", type=str, help="The path to the HDF5 file representing the DMPlex data") +parser.add_argument( + "dmplex_h5_file_path", + type=str, + help="The path to the HDF5 file representing the DMPlex data", +) args = parser.parse_args() print(f"Plotting: {args.dmplex_h5_file_path}") file_path = args.dmplex_h5_file_path dm = PETSc.DMPlex().create() -viewer = PETSc.Viewer().createHDF5(file_path, 'r') +viewer = PETSc.Viewer().createHDF5(file_path, "r") dm.load(viewer) viewer.destroy() dm.setFromOptions() @@ -25,23 +30,23 @@ # Get all edges (cone of each edge cell) edges = [] for p in range(dm.getChart()[0], dm.getChart()[1]): - #if dm.getLabelValue("celltype", p) == PETSc.DMPlex.CellType.EDGE: + # if dm.getLabelValue("celltype", p) == PETSc.DMPlex.CellType.EDGE: cone = dm.getCone(p) if len(cone) == 2: v0, v1 = cone - x0 = coords[section.getOffset(v0):section.getOffset(v0)+dim] - x1 = coords[section.getOffset(v1):section.getOffset(v1)+dim] + x0 = coords[section.getOffset(v0) : section.getOffset(v0) + dim] + x1 = coords[section.getOffset(v1) : section.getOffset(v1) + dim] edges.append((x0, x1)) # Plot using object-oriented matplotlib fig, ax = plt.subplots() for x0, x1 in edges: - ax.plot([x0[0], x1[0]], [x0[1], x1[1]], color='k', linewidth=0.5) + ax.plot([x0[0], x1[0]], [x0[1], x1[1]], color="k", linewidth=0.5) -ax.set_aspect('equal') -ax.set_title('DMPlex Mesh Edges') -ax.set_xlabel('x') -ax.set_ylabel('y') -output_path = file_path[:-3]+'.pdf' +ax.set_aspect("equal") +ax.set_title("DMPlex Mesh Edges") +ax.set_xlabel("x") +ax.set_ylabel("y") +output_path = file_path[:-3] + ".pdf" plt.savefig(output_path) print(f"Saving file to: {output_path}") diff --git a/scripts/dmplex_tools_for_particle_pusher/particle_animator.py b/scripts/dmplex_tools_for_particle_pusher/particle_animator.py index 4b4310bc6..44ccc1633 100644 --- a/scripts/dmplex_tools_for_particle_pusher/particle_animator.py +++ b/scripts/dmplex_tools_for_particle_pusher/particle_animator.py @@ -5,22 +5,37 @@ from matplotlib.animation import FuncAnimation from petsc4py import PETSc import argparse -parser = argparse.ArgumentParser(description="Animate particles moving on a DMPlex mesh.") -parser.add_argument("dmplex_h5_file_path", type=str, help="The path to the HDF5 file representing the DMPlex data") -parser.add_argument("particle_trajectory_h5_file_path", type=str, help="The path to the HDF5 file representing the particle data") + +parser = argparse.ArgumentParser( + description="Animate particles moving on a DMPlex mesh." +) +parser.add_argument( + "dmplex_h5_file_path", + type=str, + help="The path to the HDF5 file representing the DMPlex data", +) +parser.add_argument( + "particle_trajectory_h5_file_path", + type=str, + help="The path to the HDF5 file representing the particle data", +) args = parser.parse_args() -print(f"Animating particle paths from {args.particle_trajectory_h5_file_path} on DMPlex edges from {args.dmplex_h5_file_path}") +print( + f"Animating particle paths from {args.particle_trajectory_h5_file_path} on DMPlex edges from {args.dmplex_h5_file_path}" +) + def load_dmplex(file_path): dm = PETSc.DMPlex().create() - viewer = PETSc.Viewer().createHDF5(file_path, 'r') + viewer = PETSc.Viewer().createHDF5(file_path, "r") dm.load(viewer) viewer.destroy() dm.setFromOptions() dm.setUp() return dm + def get_mesh_edges(dm): # Get coordinates section = dm.getCoordinateSection() @@ -33,45 +48,49 @@ def get_mesh_edges(dm): cone = dm.getCone(p) if len(cone) == 2: v0, v1 = cone - x0 = coords[section.getOffset(v0):section.getOffset(v0)+dim] - x1 = coords[section.getOffset(v1):section.getOffset(v1)+dim] + x0 = coords[section.getOffset(v0) : section.getOffset(v0) + dim] + x1 = coords[section.getOffset(v1) : section.getOffset(v1) + dim] edges.append((x0, x1)) return edges + dm = load_dmplex(args.dmplex_h5_file_path) -#dm = load_dmplex('dmplex/expected_nonorthogonal.grd.nc.mesh.h5') +# dm = load_dmplex('dmplex/expected_nonorthogonal.grd.nc.mesh.h5') edges = get_mesh_edges(dm) # Plot using object-oriented matplotlib fig, ax = plt.subplots() + def load_particle_data(file_path): - particle_data = h5py.File(file_path,"r") + particle_data = h5py.File(file_path, "r") data_all_timesteps = list(particle_data.keys()) nstep = len(data_all_timesteps) particle_positions = [] - for it in range(0,nstep): + for it in range(0, nstep): try: group = particle_data[f"Step#{it}"] - #print(list(group.keys())) + # print(list(group.keys())) P_0 = group["POSITION_0"] P_1 = group["POSITION_1"] nparticles = len(P_0) - pdata = np.zeros((nparticles,2)) - pdata[:,0] = P_0 - pdata[:,1] = P_1 + pdata = np.zeros((nparticles, 2)) + pdata[:, 0] = P_0 + pdata[:, 1] = P_1 particle_positions.append(pdata) - except (KeyError) as error: + except KeyError as error: print(f"No particles at time step {it}: {error}") # assign empty particle data - pdata = np.empty((0,2)) + pdata = np.empty((0, 2)) particle_positions.append(pdata) return particle_positions + def update_plot(i, data, scat): scat.set_offsets(data[i]) - return scat, + return (scat,) + particle_positions = load_particle_data(args.particle_trajectory_h5_file_path) nstep = len(particle_positions) @@ -80,23 +99,25 @@ def update_plot(i, data, scat): fig, ax = plt.subplots() # Plot DMPlex edges for x0, x1 in edges: - ax.plot([x0[0], x1[0]], [x0[1], x1[1]], color='k', linewidth=0.5) + ax.plot([x0[0], x1[0]], [x0[1], x1[1]], color="k", linewidth=0.5) # Animate particles -scat = ax.scatter(particle_positions[0][:,0], particle_positions[0][:,1], c='b',s=1.0, marker='.') -ax.set_title('Particle Positions') -ax.set_xlabel('R') -ax.set_ylabel('Z') +scat = ax.scatter( + particle_positions[0][:, 0], particle_positions[0][:, 1], c="b", s=1.0, marker="." +) +ax.set_title("Particle Positions") +ax.set_xlabel("R") +ax.set_ylabel("Z") + def update(frame): - scat.set_offsets(np.c_[particle_positions[frame][:,0], particle_positions[frame][:,1]]) - return scat, + scat.set_offsets( + np.c_[particle_positions[frame][:, 0], particle_positions[frame][:, 1]] + ) + return (scat,) + ani = FuncAnimation(fig, update, frames=nstep, interval=50, blit=True) output_path = args.particle_trajectory_h5_file_path + ".animation.gif" ani.save(output_path) print(f"Saving animation of particle paths to {output_path}") -#plt.show() - - - - +# plt.show() diff --git a/scripts/dmplex_tools_for_particle_pusher/petsc_dmplex_from_hypnotoad.py b/scripts/dmplex_tools_for_particle_pusher/petsc_dmplex_from_hypnotoad.py index c1eba4661..ba2f06828 100644 --- a/scripts/dmplex_tools_for_particle_pusher/petsc_dmplex_from_hypnotoad.py +++ b/scripts/dmplex_tools_for_particle_pusher/petsc_dmplex_from_hypnotoad.py @@ -3,14 +3,21 @@ import meshio import argparse -parser = argparse.ArgumentParser(description="Create a HDF5, VTK, GMSH representation of a PETSc DMPlex mesh from a Hypnotoad grid file.") -parser.add_argument("hypnotoad_nc_file_path", type=str, help="The path to the NetCDF file containing the Hypnotoad grid data") +parser = argparse.ArgumentParser( + description="Create a HDF5, VTK, GMSH representation of a PETSc DMPlex mesh from a Hypnotoad grid file." +) +parser.add_argument( + "hypnotoad_nc_file_path", + type=str, + help="The path to the NetCDF file containing the Hypnotoad grid data", +) args = parser.parse_args() file_path = args.hypnotoad_nc_file_path print(f"Creating DMPlex from Hypnotoad grid: {file_path}") + # function to check correctly assigned labels -def check_label_value_coords(value,coords,boundary_vertex_info): +def check_label_value_coords(value, coords, boundary_vertex_info): check = False for key in boundary_vertex_info.keys(): if value == boundary_vertex_info[key]["DMFaceSetsLabel"]: @@ -20,7 +27,11 @@ def check_label_value_coords(value,coords,boundary_vertex_info): check = True break return check -def create_and_visualize_mesh(file_path,Nx,Ny,cell_list,vertex_list,boundary_vertex_info): + + +def create_and_visualize_mesh( + file_path, Nx, Ny, cell_list, vertex_list, boundary_vertex_info +): # Initialize PETSc comm = PETSc.COMM_WORLD # Create a DMPlex object @@ -28,12 +39,12 @@ def create_and_visualize_mesh(file_path,Nx,Ny,cell_list,vertex_list,boundary_ver # Define mesh parameters # for a rectangle with 2 cells and 6 vertices dim = 2 # 2D mesh - num_cells = Nx*Ny + num_cells = Nx * Ny # integers defining the cells - #cells = [[0, 1, 3, 4], [1, 2, 4, 5]] - cells = cell_list.astype('int32') - #print("cells") - #print(cells) + # cells = [[0, 1, 3, 4], [1, 2, 4, 5]] + cells = cell_list.astype("int32") + # print("cells") + # print(cells) # global list of vertices, in correct order # vertex_coords = [[0.0, 0.0], # [0.0, 1.0], @@ -51,12 +62,12 @@ def create_and_visualize_mesh(file_path,Nx,Ny,cell_list,vertex_list,boundary_ver # [2.0, 1.0], # ] vertex_coords = vertex_list - #print("vertex_coords") - #print(vertex_coords) + # print("vertex_coords") + # print(vertex_coords) # Create the mesh from the cell list dm.createFromCellList(dim, cells, vertex_coords, comm=comm) # Distribute the mesh (optional, useful for parallel runs) - #dm = dm.distribute() + # dm = dm.distribute() label_name = "Face Sets" dm.createLabel(label_name) global_label = False @@ -74,8 +85,8 @@ def create_and_visualize_mesh(file_path,Nx,Ny,cell_list,vertex_list,boundary_ver for key in boundary_vertex_info.keys(): label_value = boundary_vertex_info[key]["DMFaceSetsLabel"] boundary_vertex_indices = boundary_vertex_info[key]["ivertex"] - #print(key) - #print(boundary_vertex_indices) + # print(key) + # print(boundary_vertex_indices) # loop over faces for face in range(face_start, face_end): # vertices supporting this face @@ -83,7 +94,7 @@ def create_and_visualize_mesh(file_path,Nx,Ny,cell_list,vertex_list,boundary_ver # n.b. need to subtract vertex_start # to get back the input global index # if all points in cone are in the boundary, this is a boundary face - if all( (v-vertex_start) in boundary_vertex_indices for v in cone): + if all((v - vertex_start) in boundary_vertex_indices for v in cone): dm.setLabelValue("Face Sets", face, label_value) # Print labeled boundary faces @@ -95,41 +106,54 @@ def create_and_visualize_mesh(file_path,Nx,Ny,cell_list,vertex_list,boundary_ver value = dm.getLabelValue(label_name, face) cone = dm.getCone(face) coords = vertex_coords[cone - vertex_start] - if value in [100,200,300,400]: - #print(f"Face {face} labeled with value {value}") - #print(f"Face {face} coords {coords[0]}") - check = check_label_value_coords(value,coords[0],boundary_vertex_info) - #print(f"Test passed: {check}") - mesh_name = file_path+".mesh" - mesh_path = mesh_name+".vtk" + if value in [100, 200, 300, 400]: + # print(f"Face {face} labeled with value {value}") + # print(f"Face {face} coords {coords[0]}") + check = check_label_value_coords(value, coords[0], boundary_vertex_info) + # print(f"Test passed: {check}") + mesh_name = file_path + ".mesh" + mesh_path = mesh_name + ".vtk" # Write the mesh to a VTK file - viewer = PETSc.Viewer().createVTK(mesh_path, mode=PETSc.Viewer.Mode.WRITE, comm=comm) + viewer = PETSc.Viewer().createVTK( + mesh_path, mode=PETSc.Viewer.Mode.WRITE, comm=comm + ) dm.view(viewer) viewer.destroy() print(f"Saved DMPlex view: {mesh_path}") # Write the mesh to a .h5 file - mesh_path = mesh_name+".h5" - #viewer = PETSc.Viewer().createHDF5('mesh.h5', mode=PETSc.Viewer.Mode.WRITE, comm=comm) - viewer = PETSc.ViewerHDF5().create(mesh_path, mode=PETSc.Viewer.Mode.WRITE, comm=comm) + mesh_path = mesh_name + ".h5" + # viewer = PETSc.Viewer().createHDF5('mesh.h5', mode=PETSc.Viewer.Mode.WRITE, comm=comm) + viewer = PETSc.ViewerHDF5().create( + mesh_path, mode=PETSc.Viewer.Mode.WRITE, comm=comm + ) dm.view(viewer) viewer.destroy() print(f"Saved DMPlex view: {mesh_path}") # Read the mesh from .h5 dmtest = PETSc.DMPlex().create(comm=comm) - viewer = PETSc.ViewerHDF5().create(mesh_path, mode=PETSc.Viewer.Mode.READ, comm=comm) + viewer = PETSc.ViewerHDF5().create( + mesh_path, mode=PETSc.Viewer.Mode.READ, comm=comm + ) dmtest.load(viewer) viewer.destroy() # Write the reloaded mesh to a VTK file - mesh_path = mesh_name+"_h5_to_vtk.vtk" - viewer = PETSc.Viewer().createVTK(mesh_path, mode=PETSc.Viewer.Mode.WRITE, comm=comm) + mesh_path = mesh_name + "_h5_to_vtk.vtk" + viewer = PETSc.Viewer().createVTK( + mesh_path, mode=PETSc.Viewer.Mode.WRITE, comm=comm + ) dmtest.view(viewer) viewer.destroy() # write mesh to gmsh format with ascii output, for use in NESO-particles tests - mesh = meshio.read(mesh_name+".vtk") - mesh_path = mesh_name+".msh" - mesh.write(mesh_path,binary=False) + mesh = meshio.read(mesh_name + ".vtk") + mesh_path = mesh_name + ".msh" + mesh.write(mesh_path, binary=False) print(f"Saved DMPlex view: {mesh_path}") return None -Nx,Ny,cell_vertices,vertex_list,boundary_vertex_info = plot_corners_get_dmplex_data(file_path) -create_and_visualize_mesh(file_path,Nx,Ny,cell_vertices,vertex_list,boundary_vertex_info) + +Nx, Ny, cell_vertices, vertex_list, boundary_vertex_info = plot_corners_get_dmplex_data( + file_path +) +create_and_visualize_mesh( + file_path, Nx, Ny, cell_vertices, vertex_list, boundary_vertex_info +) diff --git a/scripts/dmplex_tools_for_particle_pusher/plot_corners_functions.py b/scripts/dmplex_tools_for_particle_pusher/plot_corners_functions.py index 5d15b69d4..d9e599f8f 100644 --- a/scripts/dmplex_tools_for_particle_pusher/plot_corners_functions.py +++ b/scripts/dmplex_tools_for_particle_pusher/plot_corners_functions.py @@ -4,7 +4,10 @@ import os import subprocess -def save_ivertex_indices_to_netcdf(source_file,destination_file,Rpoints_full,Zpoints_full): + +def save_ivertex_indices_to_netcdf( + source_file, destination_file, Rpoints_full, Zpoints_full +): # copy source file to destination os.system(f"cp {source_file} {destination_file}") # open destination file @@ -13,120 +16,145 @@ def save_ivertex_indices_to_netcdf(source_file,destination_file,Rpoints_full,Zpo # get data including y guards # but note that Rpoints_full,Zpoints_full created # without y guards as the DMPlex should not have guards - data_Rxy = np.copy(dataset.variables['Rxy_corners'][:]) - data_Zxy = np.copy(dataset.variables['Zxy_corners'][:]) - data_Rxy_lr = np.copy(dataset.variables['Rxy_lower_right_corners'][:]) - data_Zxy_lr = np.copy(dataset.variables['Zxy_lower_right_corners'][:]) - data_Rxy_ur = np.copy(dataset.variables['Rxy_upper_right_corners'][:]) - data_Zxy_ur = np.copy(dataset.variables['Zxy_upper_right_corners'][:]) - data_Rxy_ul = np.copy(dataset.variables['Rxy_upper_left_corners'][:]) - data_Zxy_ul = np.copy(dataset.variables['Zxy_upper_left_corners'][:]) - - ivertex_corners = get_boutxx_corner_index(data_Rxy,data_Zxy,Rpoints_full,Zpoints_full) - ivertex_corners_lr = get_boutxx_corner_index(data_Rxy_lr,data_Zxy_lr,Rpoints_full,Zpoints_full) - ivertex_corners_ul = get_boutxx_corner_index(data_Rxy_ul,data_Zxy_ul,Rpoints_full,Zpoints_full) - ivertex_corners_ur = get_boutxx_corner_index(data_Rxy_ur,data_Zxy_ur,Rpoints_full,Zpoints_full) + data_Rxy = np.copy(dataset.variables["Rxy_corners"][:]) + data_Zxy = np.copy(dataset.variables["Zxy_corners"][:]) + data_Rxy_lr = np.copy(dataset.variables["Rxy_lower_right_corners"][:]) + data_Zxy_lr = np.copy(dataset.variables["Zxy_lower_right_corners"][:]) + data_Rxy_ur = np.copy(dataset.variables["Rxy_upper_right_corners"][:]) + data_Zxy_ur = np.copy(dataset.variables["Zxy_upper_right_corners"][:]) + data_Rxy_ul = np.copy(dataset.variables["Rxy_upper_left_corners"][:]) + data_Zxy_ul = np.copy(dataset.variables["Zxy_upper_left_corners"][:]) + + ivertex_corners = get_boutxx_corner_index( + data_Rxy, data_Zxy, Rpoints_full, Zpoints_full + ) + ivertex_corners_lr = get_boutxx_corner_index( + data_Rxy_lr, data_Zxy_lr, Rpoints_full, Zpoints_full + ) + ivertex_corners_ul = get_boutxx_corner_index( + data_Rxy_ul, data_Zxy_ul, Rpoints_full, Zpoints_full + ) + ivertex_corners_ur = get_boutxx_corner_index( + data_Rxy_ur, data_Zxy_ur, Rpoints_full, Zpoints_full + ) # index saved as double as BOUT++ can only handle double Field2D - ptr = dataset.createVariable("ivertex_lower_left_corners","f8",("x", "y")) - ptr.setncattr('bout_type','Field2D') + ptr = dataset.createVariable("ivertex_lower_left_corners", "f8", ("x", "y")) + ptr.setncattr("bout_type", "Field2D") ptr[:] = ivertex_corners - ptr = dataset.createVariable("ivertex_lower_right_corners","f8",("x", "y")) - ptr.setncattr('bout_type','Field2D') + ptr = dataset.createVariable("ivertex_lower_right_corners", "f8", ("x", "y")) + ptr.setncattr("bout_type", "Field2D") ptr[:] = ivertex_corners_lr - ptr = dataset.createVariable("ivertex_upper_right_corners","f8",("x", "y")) - ptr.setncattr('bout_type','Field2D') + ptr = dataset.createVariable("ivertex_upper_right_corners", "f8", ("x", "y")) + ptr.setncattr("bout_type", "Field2D") ptr[:] = ivertex_corners_ur - ptr = dataset.createVariable("ivertex_upper_left_corners","f8",("x", "y")) - ptr.setncattr('bout_type','Field2D') + ptr = dataset.createVariable("ivertex_upper_left_corners", "f8", ("x", "y")) + ptr.setncattr("bout_type", "Field2D") ptr[:] = ivertex_corners_ul # add a new dimension to the Hypnotoad output file for the global vertex list - dataset.createDimension('nvertices', len(Rpoints_full)) + dataset.createDimension("nvertices", len(Rpoints_full)) # add the global vertex lists - ptr = dataset.createVariable("global_vertex_list_R","f8",("nvertices",)) + ptr = dataset.createVariable("global_vertex_list_R", "f8", ("nvertices",)) ptr[:] = Rpoints_full ptr.units = "length" ptr.description = "global list of R points of vertices in a hypnotoad mesh" ptr.source = "generated in python external to hypnotoad" - ptr = dataset.createVariable("global_vertex_list_Z","f8",("nvertices",)) + ptr = dataset.createVariable("global_vertex_list_Z", "f8", ("nvertices",)) ptr.units = "length" ptr.description = "global list of Z points of vertices in a hypnotoad mesh" ptr.source = "generated in python external to hypnotoad" ptr[:] = Zpoints_full # record version information - dataset.setncattr("vertices_calculation_version_git_hash", subprocess.check_output('git rev-parse HEAD',shell=True).decode('utf-8').strip()) + dataset.setncattr( + "vertices_calculation_version_git_hash", + subprocess.check_output("git rev-parse HEAD", shell=True) + .decode("utf-8") + .strip(), + ) dataset.setncattr("original_hypnotoad_file_name", os.path.basename(source_file)) dataset.close() print(f"Saving modified Hypnotoad mesh file to {destination_file}") return None -def isapprox(a,b,tol=1.0e-8): - if abs(a - b) < tol: - return True - else: - return False -def unique_points_2D(corners_Rxy,corners_Zxy): +def isapprox(a, b, tol=1.0e-8): + if abs(a - b) < tol: + return True + else: + return False + + +def unique_points_2D(corners_Rxy, corners_Zxy): Nx, Ny = corners_Rxy.shape - unique=True - for ix in range(0,Nx): - for iy in range(0,Ny): - for ixp in range(0,Nx): - for iyp in range(0,Ny): + unique = True + for ix in range(0, Nx): + for iy in range(0, Ny): + for ixp in range(0, Nx): + for iyp in range(0, Ny): if ix == ixp and iy == iyp: continue - if (isapprox(corners_Rxy[ix,iy],corners_Rxy[ixp,iyp]) and isapprox(corners_Zxy[ix,iy],corners_Zxy[ixp,iyp])): - ic = iy + Ny*ix - #print("non-unique point:",ix,iy,ic) + if isapprox( + corners_Rxy[ix, iy], corners_Rxy[ixp, iyp] + ) and isapprox(corners_Zxy[ix, iy], corners_Zxy[ixp, iyp]): + ic = iy + Ny * ix + # print("non-unique point:",ix,iy,ic) unique = False - #print("unique points =",unique) + # print("unique points =",unique) return None -def unique_points_1D(corners_Rxy,corners_Zxy): - Nc, = corners_Rxy.shape - unique=True - for ic in range(0,Nc): - for icp in range(0,Nc): + +def unique_points_1D(corners_Rxy, corners_Zxy): + (Nc,) = corners_Rxy.shape + unique = True + for ic in range(0, Nc): + for icp in range(0, Nc): if ic == icp: continue - if (isapprox(corners_Rxy[ic],corners_Rxy[icp]) and isapprox(corners_Zxy[ic],corners_Zxy[icp])): - #print("non-unique point:",ix,iy,ic) + if isapprox(corners_Rxy[ic], corners_Rxy[icp]) and isapprox( + corners_Zxy[ic], corners_Zxy[icp] + ): + # print("non-unique point:",ix,iy,ic) unique = False - #print("unique points =",unique) + # print("unique points =",unique) return None -def remove_nonunique_points(corners_Rxy,corners_Zxy): - Nc, = corners_Rxy.shape + +def remove_nonunique_points(corners_Rxy, corners_Zxy): + (Nc,) = corners_Rxy.shape duplicate_indices = [] - for ic in range(0,Nc): - for icp in range(0,Nc): + for ic in range(0, Nc): + for icp in range(0, Nc): if ic == icp: continue - if (isapprox(corners_Rxy[ic],corners_Rxy[icp]) and isapprox(corners_Zxy[ic],corners_Zxy[icp])): - #print("non-unique point:",icp) + if isapprox(corners_Rxy[ic], corners_Rxy[icp]) and isapprox( + corners_Zxy[ic], corners_Zxy[icp] + ): + # print("non-unique point:",icp) # store the duplicate index, if we have not already stored the index for this point if (ic not in duplicate_indices) and (icp not in duplicate_indices): duplicate_indices.append(icp) - #print(duplicate_indices) + # print(duplicate_indices) if len(duplicate_indices) > 0: - corners_Rxy = np.delete(corners_Rxy,duplicate_indices) - corners_Zxy = np.delete(corners_Zxy,duplicate_indices) + corners_Rxy = np.delete(corners_Rxy, duplicate_indices) + corners_Zxy = np.delete(corners_Zxy, duplicate_indices) return corners_Rxy, corners_Zxy + # identify which vertex an (R,Z) pair, or print error message -def convert_R_Z_to_vertex_index(R,Z,R_vertices,Z_vertices,failure=None): +def convert_R_Z_to_vertex_index(R, Z, R_vertices, Z_vertices, failure=None): Nc = len(R_vertices) - for ic in range(0,Nc): - if isapprox(R,R_vertices[ic]) and isapprox(Z,Z_vertices[ic]): + for ic in range(0, Nc): + if isapprox(R, R_vertices[ic]) and isapprox(Z, Z_vertices[ic]): ic_vertex = ic return ic_vertex - #print("Failed to find (R,Z) index") + # print("Failed to find (R,Z) index") return failure -def get_cell_vertex_list_inner(R_ll,Z_ll, R_lr, Z_lr, - R_ur, Z_ur, R_ul, Z_ul, - R_vertices, Z_vertices): + +def get_cell_vertex_list_inner( + R_ll, Z_ll, R_lr, Z_lr, R_ur, Z_ur, R_ul, Z_ul, R_vertices, Z_vertices +): # R coords defining the vertcies R_local_vertex_list = np.array([R_ll, R_lr, R_ur, R_ul]) # R location of cell centre @@ -139,10 +167,10 @@ def get_cell_vertex_list_inner(R_ll,Z_ll, R_lr, Z_lr, # midpoint of the cell, measured # with respect to the major radial direction theta_list = np.zeros(4) - for j in range(0,4): + for j in range(0, 4): R = R_local_vertex_list[j] - R_mid Z = Z_local_vertex_list[j] - Z_mid - theta_list[j] = np.arctan2(Z,R) + theta_list[j] = np.arctan2(Z, R) # order theta in increasing order # to get the order for anticlockwise # specification of the vertices @@ -151,64 +179,89 @@ def get_cell_vertex_list_inner(R_ll,Z_ll, R_lr, Z_lr, Z_local_vertex_sorted = Z_local_vertex_list[sort_indices] vertices = [] # append the vertices in anti-clockwise order - for j in range(0,4): + for j in range(0, 4): R = R_local_vertex_sorted[j] Z = Z_local_vertex_sorted[j] - vertices.append(convert_R_Z_to_vertex_index(R,Z,R_vertices,Z_vertices)) + vertices.append(convert_R_Z_to_vertex_index(R, Z, R_vertices, Z_vertices)) return vertices -def get_cell_vertex_list(Rxy_ll, Zxy_ll, - Rxy_lr, Zxy_lr, Rxy_ur, Zxy_ur, Rxy_ul, Zxy_ul, - R_vertices, Z_vertices): + +def get_cell_vertex_list( + Rxy_ll, + Zxy_ll, + Rxy_lr, + Zxy_lr, + Rxy_ur, + Zxy_ur, + Rxy_ul, + Zxy_ul, + R_vertices, + Z_vertices, +): Nx, Ny = np.shape(Rxy_ll) - Ncells = Nx*Ny - cell_vertices = np.zeros((Ncells,4),dtype=int) - cell_vertices_RZ = np.zeros((Ncells,4,2),dtype=float) - cell_vertices[:,:] = -3000 - for ix in range(0,Nx): - for iy in range(0,Ny): - icell = iy + Ny*ix - R_ll = Rxy_ll[ix,iy] - Z_ll = Zxy_ll[ix,iy] - R_lr = Rxy_lr[ix,iy] - Z_lr = Zxy_lr[ix,iy] - R_ur = Rxy_ur[ix,iy] - Z_ur = Zxy_ur[ix,iy] - R_ul = Rxy_ul[ix,iy] - Z_ul = Zxy_ul[ix,iy] - cell_vertices[icell,:] = get_cell_vertex_list_inner(R_ll,Z_ll, R_lr, Z_lr, - R_ur, Z_ur, R_ul, Z_ul, - R_vertices, Z_vertices) + Ncells = Nx * Ny + cell_vertices = np.zeros((Ncells, 4), dtype=int) + cell_vertices_RZ = np.zeros((Ncells, 4, 2), dtype=float) + cell_vertices[:, :] = -3000 + for ix in range(0, Nx): + for iy in range(0, Ny): + icell = iy + Ny * ix + R_ll = Rxy_ll[ix, iy] + Z_ll = Zxy_ll[ix, iy] + R_lr = Rxy_lr[ix, iy] + Z_lr = Zxy_lr[ix, iy] + R_ur = Rxy_ur[ix, iy] + Z_ur = Zxy_ur[ix, iy] + R_ul = Rxy_ul[ix, iy] + Z_ul = Zxy_ul[ix, iy] + cell_vertices[icell, :] = get_cell_vertex_list_inner( + R_ll, Z_ll, R_lr, Z_lr, R_ur, Z_ur, R_ul, Z_ul, R_vertices, Z_vertices + ) # this list must be in the same order as the list of cell vertices above - cell_vertices_RZ[icell,0,:] = [R_ll,Z_ll] - cell_vertices_RZ[icell,1,:] = [R_lr,Z_lr] - cell_vertices_RZ[icell,2,:] = [R_ur,Z_ur] - cell_vertices_RZ[icell,3,:] = [R_ul,Z_ul] + cell_vertices_RZ[icell, 0, :] = [R_ll, Z_ll] + cell_vertices_RZ[icell, 1, :] = [R_lr, Z_lr] + cell_vertices_RZ[icell, 2, :] = [R_ur, Z_ur] + cell_vertices_RZ[icell, 3, :] = [R_ul, Z_ul] # get a global list of vertices - Nc, = np.shape(R_vertices) - vertex_list = np.zeros((Nc,2),dtype=float) - for ic in range(0,Nc): - vertex_list[ic,0] = R_vertices[ic] - vertex_list[ic,1] = Z_vertices[ic] + (Nc,) = np.shape(R_vertices) + vertex_list = np.zeros((Nc, 2), dtype=float) + for ic in range(0, Nc): + vertex_list[ic, 0] = R_vertices[ic] + vertex_list[ic, 1] = Z_vertices[ic] return cell_vertices, cell_vertices_RZ, vertex_list -def get_boutxx_corner_index(Rxy_corners,Zxy_corners,R_vertices,Z_vertices): + +def get_boutxx_corner_index(Rxy_corners, Zxy_corners, R_vertices, Z_vertices): Nx, Ny = Rxy_corners.shape - iglobal_corners = np.zeros((Nx,Ny),dtype=int) - if Nx*Ny > len(R_vertices): + iglobal_corners = np.zeros((Nx, Ny), dtype=int) + if Nx * Ny > len(R_vertices): failure_return = -1 else: failure_return = None - for ix in range(0,Nx): - for iy in range(0,Ny): + for ix in range(0, Nx): + for iy in range(0, Ny): # find the global index for this corner, and store - iglobal_corners[ix,iy] = convert_R_Z_to_vertex_index(Rxy_corners[ix,iy],Zxy_corners[ix,iy], - R_vertices,Z_vertices,failure=failure_return) + iglobal_corners[ix, iy] = convert_R_Z_to_vertex_index( + Rxy_corners[ix, iy], + Zxy_corners[ix, iy], + R_vertices, + Z_vertices, + failure=failure_return, + ) return iglobal_corners -def get_pfr_lower_boundary_vertices(ivertex_corners,data_Rxy,data_Zxy, - ivertex_corners_ul,data_Rxy_ul,data_Zxy_ul, - y_boundary_guards,jyseps1_1,exclude_y_guard_cells=False): + +def get_pfr_lower_boundary_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + ivertex_corners_ul, + data_Rxy_ul, + data_Zxy_ul, + y_boundary_guards, + jyseps1_1, + exclude_y_guard_cells=False, +): Nx, Ny = ivertex_corners.shape ivertex_pfr_lower = [] Rxy_pfr_lower = [] @@ -217,26 +270,37 @@ def get_pfr_lower_boundary_vertices(ivertex_corners,data_Rxy,data_Zxy, jy_bndry = 0 else: jy_bndry = y_boundary_guards - lim1 = jyseps1_1+jy_bndry+1 - for j in range(jy_bndry,lim1): - ivertex_pfr_lower.append(ivertex_corners[0,j]) - Rxy_pfr_lower.append(data_Rxy[0,j]) - Zxy_pfr_lower.append(data_Zxy[0,j]) + lim1 = jyseps1_1 + jy_bndry + 1 + for j in range(jy_bndry, lim1): + ivertex_pfr_lower.append(ivertex_corners[0, j]) + Rxy_pfr_lower.append(data_Rxy[0, j]) + Zxy_pfr_lower.append(data_Zxy[0, j]) lim2 = Ny - jy_bndry - jyseps1_1 - 1 lim3 = Ny - jy_bndry - for j in range(lim2,lim3): - ivertex_pfr_lower.append(ivertex_corners[0,j]) - Rxy_pfr_lower.append(data_Rxy[0,j]) - Zxy_pfr_lower.append(data_Zxy[0,j]) - ivertex_pfr_lower.append(ivertex_corners_ul[0,lim3-1]) - Rxy_pfr_lower.append(data_Rxy_ul[0,lim3-1]) - Zxy_pfr_lower.append(data_Zxy_ul[0,lim3-1]) + for j in range(lim2, lim3): + ivertex_pfr_lower.append(ivertex_corners[0, j]) + Rxy_pfr_lower.append(data_Rxy[0, j]) + Zxy_pfr_lower.append(data_Zxy[0, j]) + ivertex_pfr_lower.append(ivertex_corners_ul[0, lim3 - 1]) + Rxy_pfr_lower.append(data_Rxy_ul[0, lim3 - 1]) + Zxy_pfr_lower.append(data_Zxy_ul[0, lim3 - 1]) return ivertex_pfr_lower, Rxy_pfr_lower, Zxy_pfr_lower -def get_pfr_upper_boundary_vertices(ivertex_corners,data_Rxy,data_Zxy, - ivertex_corners_ul,data_Rxy_ul,data_Zxy_ul, - y_boundary_guards,jyseps1_2,jyseps2_1,jyseps2_2, - ny_inner, exclude_y_guard_cells=False): + +def get_pfr_upper_boundary_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + ivertex_corners_ul, + data_Rxy_ul, + data_Zxy_ul, + y_boundary_guards, + jyseps1_2, + jyseps2_1, + jyseps2_2, + ny_inner, + exclude_y_guard_cells=False, +): ivertex_pfr_upper = [] Rxy_pfr_upper = [] Zxy_pfr_upper = [] @@ -244,57 +308,82 @@ def get_pfr_upper_boundary_vertices(ivertex_corners,data_Rxy,data_Zxy, jyseps2_1g = jyseps2_1 else: jyseps2_1g = jyseps2_1 + y_boundary_guards - for j in range(ny_inner,jyseps1_2+1): - ivertex_pfr_upper.append(ivertex_corners[0,j]) - Rxy_pfr_upper.append(data_Rxy[0,j]) - Zxy_pfr_upper.append(data_Zxy[0,j]) - for j in range(jyseps2_1+1,ny_inner): - ivertex_pfr_upper.append(ivertex_corners[0,j]) - Rxy_pfr_upper.append(data_Rxy[0,j]) - Zxy_pfr_upper.append(data_Zxy[0,j]) - j = ny_inner-1 - ivertex_pfr_upper.append(ivertex_corners_ul[0,j]) - Rxy_pfr_upper.append(data_Rxy_ul[0,j]) - Zxy_pfr_upper.append(data_Zxy_ul[0,j]) + for j in range(ny_inner, jyseps1_2 + 1): + ivertex_pfr_upper.append(ivertex_corners[0, j]) + Rxy_pfr_upper.append(data_Rxy[0, j]) + Zxy_pfr_upper.append(data_Zxy[0, j]) + for j in range(jyseps2_1 + 1, ny_inner): + ivertex_pfr_upper.append(ivertex_corners[0, j]) + Rxy_pfr_upper.append(data_Rxy[0, j]) + Zxy_pfr_upper.append(data_Zxy[0, j]) + j = ny_inner - 1 + ivertex_pfr_upper.append(ivertex_corners_ul[0, j]) + Rxy_pfr_upper.append(data_Rxy_ul[0, j]) + Zxy_pfr_upper.append(data_Zxy_ul[0, j]) return ivertex_pfr_upper, Rxy_pfr_upper, Zxy_pfr_upper -def get_vac_left_boundary_vertices(ivertex_corners_lr,data_Rxy_lr,data_Zxy_lr, - ivertex_corners_ur,data_Rxy_ur,data_Zxy_ur, - y_boundary_guards,ny_inner): + +def get_vac_left_boundary_vertices( + ivertex_corners_lr, + data_Rxy_lr, + data_Zxy_lr, + ivertex_corners_ur, + data_Rxy_ur, + data_Zxy_ur, + y_boundary_guards, + ny_inner, +): ivertex_sol_vac_left = [] Rxy_sol_vac_left = [] Zxy_sol_vac_left = [] - jlim = ny_inner#+2*y_boundary_guards - for j in range(0,jlim): - ivertex_sol_vac_left.append(ivertex_corners_lr[-1,j]) - Rxy_sol_vac_left.append(data_Rxy_lr[-1,j]) - Zxy_sol_vac_left.append(data_Zxy_lr[-1,j]) + jlim = ny_inner # +2*y_boundary_guards + for j in range(0, jlim): + ivertex_sol_vac_left.append(ivertex_corners_lr[-1, j]) + Rxy_sol_vac_left.append(data_Rxy_lr[-1, j]) + Zxy_sol_vac_left.append(data_Zxy_lr[-1, j]) j = jlim - 1 - ivertex_sol_vac_left.append(ivertex_corners_ur[-1,j]) - Rxy_sol_vac_left.append(data_Rxy_ur[-1,j]) - Zxy_sol_vac_left.append(data_Zxy_ur[-1,j]) + ivertex_sol_vac_left.append(ivertex_corners_ur[-1, j]) + Rxy_sol_vac_left.append(data_Rxy_ur[-1, j]) + Zxy_sol_vac_left.append(data_Zxy_ur[-1, j]) return ivertex_sol_vac_left, Rxy_sol_vac_left, Zxy_sol_vac_left -def get_vac_right_boundary_vertices(ivertex_corners_lr,data_Rxy_lr,data_Zxy_lr, - ivertex_corners_ur,data_Rxy_ur,data_Zxy_ur, - jyseps1_2,ny_inner): + +def get_vac_right_boundary_vertices( + ivertex_corners_lr, + data_Rxy_lr, + data_Zxy_lr, + ivertex_corners_ur, + data_Rxy_ur, + data_Zxy_ur, + jyseps1_2, + ny_inner, +): Nx, Ny = ivertex_corners_lr.shape ivertex_sol_vac_right = [] Rxy_sol_vac_right = [] Zxy_sol_vac_right = [] - j = ny_inner#jyseps1_2+1 - ivertex_sol_vac_right.append(ivertex_corners_lr[-1,j]) - Rxy_sol_vac_right.append(data_Rxy_lr[-1,j]) - Zxy_sol_vac_right.append(data_Zxy_lr[-1,j]) - for j in range(ny_inner,Ny):#jyseps1_2+1 - ivertex_sol_vac_right.append(ivertex_corners_ur[-1,j]) - Rxy_sol_vac_right.append(data_Rxy_ur[-1,j]) - Zxy_sol_vac_right.append(data_Zxy_ur[-1,j]) + j = ny_inner # jyseps1_2+1 + ivertex_sol_vac_right.append(ivertex_corners_lr[-1, j]) + Rxy_sol_vac_right.append(data_Rxy_lr[-1, j]) + Zxy_sol_vac_right.append(data_Zxy_lr[-1, j]) + for j in range(ny_inner, Ny): # jyseps1_2+1 + ivertex_sol_vac_right.append(ivertex_corners_ur[-1, j]) + Rxy_sol_vac_right.append(data_Rxy_ur[-1, j]) + Zxy_sol_vac_right.append(data_Zxy_ur[-1, j]) return ivertex_sol_vac_right, Rxy_sol_vac_right, Zxy_sol_vac_right -def get_core_boundary_vertices(ivertex_corners, data_Rxy, data_Zxy, - jyseps1_1, jyseps2_1, jyseps1_2, jyseps2_2, - y_boundary_guards, exclude_y_guard_cells=False): + +def get_core_boundary_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + jyseps1_1, + jyseps2_1, + jyseps1_2, + jyseps2_2, + y_boundary_guards, + exclude_y_guard_cells=False, +): ivertex_core = [] Rxy_core = [] Zxy_core = [] @@ -304,97 +393,123 @@ def get_core_boundary_vertices(ivertex_corners, data_Rxy, data_Zxy, else: jyseps1_1g = jyseps1_1 + y_boundary_guards jyseps2_1g = jyseps2_1 + y_boundary_guards - for j in range(jyseps1_1g+1,jyseps2_1g+1): - ivertex_core.append(ivertex_corners[0,j]) - Rxy_core.append(data_Rxy[0,j]) - Zxy_core.append(data_Zxy[0,j]) - for j in range(jyseps1_2+1,jyseps2_2+1): - ivertex_core.append(ivertex_corners[0,j]) - Rxy_core.append(data_Rxy[0,j]) - Zxy_core.append(data_Zxy[0,j]) + for j in range(jyseps1_1g + 1, jyseps2_1g + 1): + ivertex_core.append(ivertex_corners[0, j]) + Rxy_core.append(data_Rxy[0, j]) + Zxy_core.append(data_Zxy[0, j]) + for j in range(jyseps1_2 + 1, jyseps2_2 + 1): + ivertex_core.append(ivertex_corners[0, j]) + Rxy_core.append(data_Rxy[0, j]) + Zxy_core.append(data_Zxy[0, j]) return ivertex_core, Rxy_core, Zxy_core -def get_target_ll_vertices(ivertex_corners, data_Rxy, data_Zxy, - ivertex_corners_lr, data_Rxy_lr, data_Zxy_lr): + +def get_target_ll_vertices( + ivertex_corners, data_Rxy, data_Zxy, ivertex_corners_lr, data_Rxy_lr, data_Zxy_lr +): ivertex_target_ll = [] Rxy_target_ll = [] Zxy_target_ll = [] Nx, Ny = ivertex_corners.shape j = 0 - for i in range(0,Nx): - ivertex_target_ll.append(ivertex_corners[i,j]) - Rxy_target_ll.append(data_Rxy[i,j]) - Zxy_target_ll.append(data_Zxy[i,j]) + for i in range(0, Nx): + ivertex_target_ll.append(ivertex_corners[i, j]) + Rxy_target_ll.append(data_Rxy[i, j]) + Zxy_target_ll.append(data_Zxy[i, j]) i = Nx - 1 - ivertex_target_ll.append(ivertex_corners_lr[i,j]) - Rxy_target_ll.append(data_Rxy_lr[i,j]) - Zxy_target_ll.append(data_Zxy_lr[i,j]) + ivertex_target_ll.append(ivertex_corners_lr[i, j]) + Rxy_target_ll.append(data_Rxy_lr[i, j]) + Zxy_target_ll.append(data_Zxy_lr[i, j]) return ivertex_target_ll, Rxy_target_ll, Zxy_target_ll -def get_target_ul_vertices(ivertex_corners_ul, data_Rxy_ul, data_Zxy_ul, - ivertex_corners_ur, data_Rxy_ur, data_Zxy_ur, - ny_inner): + +def get_target_ul_vertices( + ivertex_corners_ul, + data_Rxy_ul, + data_Zxy_ul, + ivertex_corners_ur, + data_Rxy_ur, + data_Zxy_ur, + ny_inner, +): ivertex_target_ul = [] Rxy_target_ul = [] Zxy_target_ul = [] Nx, Ny = ivertex_corners_ul.shape j = ny_inner - 1 - for i in range(0,Nx): - ivertex_target_ul.append(ivertex_corners_ul[i,j]) - Rxy_target_ul.append(data_Rxy_ul[i,j]) - Zxy_target_ul.append(data_Zxy_ul[i,j]) + for i in range(0, Nx): + ivertex_target_ul.append(ivertex_corners_ul[i, j]) + Rxy_target_ul.append(data_Rxy_ul[i, j]) + Zxy_target_ul.append(data_Zxy_ul[i, j]) i = Nx - 1 - ivertex_target_ul.append(ivertex_corners_ur[i,j]) - Rxy_target_ul.append(data_Rxy_ur[i,j]) - Zxy_target_ul.append(data_Zxy_ur[i,j]) + ivertex_target_ul.append(ivertex_corners_ur[i, j]) + Rxy_target_ul.append(data_Rxy_ur[i, j]) + Zxy_target_ul.append(data_Zxy_ur[i, j]) return ivertex_target_ul, Rxy_target_ul, Zxy_target_ul -def get_target_ur_vertices(ivertex_corners, data_Rxy, data_Zxy, - ivertex_corners_lr, data_Rxy_lr, data_Zxy_lr, - ny_inner): + +def get_target_ur_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + ivertex_corners_lr, + data_Rxy_lr, + data_Zxy_lr, + ny_inner, +): ivertex_target_ur = [] Rxy_target_ur = [] Zxy_target_ur = [] Nx, Ny = ivertex_corners.shape j = ny_inner - for i in range(0,Nx): - ivertex_target_ur.append(ivertex_corners[i,j]) - Rxy_target_ur.append(data_Rxy[i,j]) - Zxy_target_ur.append(data_Zxy[i,j]) + for i in range(0, Nx): + ivertex_target_ur.append(ivertex_corners[i, j]) + Rxy_target_ur.append(data_Rxy[i, j]) + Zxy_target_ur.append(data_Zxy[i, j]) i = Nx - 1 - ivertex_target_ur.append(ivertex_corners_lr[i,j]) - Rxy_target_ur.append(data_Rxy_lr[i,j]) - Zxy_target_ur.append(data_Zxy_lr[i,j]) + ivertex_target_ur.append(ivertex_corners_lr[i, j]) + Rxy_target_ur.append(data_Rxy_lr[i, j]) + Zxy_target_ur.append(data_Zxy_lr[i, j]) return ivertex_target_ur, Rxy_target_ur, Zxy_target_ur -def get_target_lr_vertices(ivertex_corners_ur, data_Rxy_ur, data_Zxy_ur, - ivertex_corners_ul, data_Rxy_ul, data_Zxy_ul): + +def get_target_lr_vertices( + ivertex_corners_ur, + data_Rxy_ur, + data_Zxy_ur, + ivertex_corners_ul, + data_Rxy_ul, + data_Zxy_ul, +): ivertex_target_lr = [] Rxy_target_lr = [] Zxy_target_lr = [] Nx, Ny = ivertex_corners_ur.shape j = Ny - 1 - for i in range(0,Nx): - ivertex_target_lr.append(ivertex_corners_ur[i,j]) - Rxy_target_lr.append(data_Rxy_ur[i,j]) - Zxy_target_lr.append(data_Zxy_ur[i,j]) + for i in range(0, Nx): + ivertex_target_lr.append(ivertex_corners_ur[i, j]) + Rxy_target_lr.append(data_Rxy_ur[i, j]) + Zxy_target_lr.append(data_Zxy_ur[i, j]) i = 0 - ivertex_target_lr.append(ivertex_corners_ul[i,j]) - Rxy_target_lr.append(data_Rxy_ul[i,j]) - Zxy_target_lr.append(data_Zxy_ul[i,j]) + ivertex_target_lr.append(ivertex_corners_ul[i, j]) + Rxy_target_lr.append(data_Rxy_ul[i, j]) + Zxy_target_lr.append(data_Zxy_ul[i, j]) return ivertex_target_lr, Rxy_target_lr, Zxy_target_lr -def plot_corners_get_dmplex_data(file_path,interactive_plot=False,print_cells_to_screen_output=False): + +def plot_corners_get_dmplex_data( + file_path, interactive_plot=False, print_cells_to_screen_output=False +): dataset = nc.Dataset(file_path) - y_boundary_guards = np.copy(dataset.variables['y_boundary_guards'][...]) - ixseps1 = np.copy(dataset.variables['ixseps1'][...]) - ixseps2 = np.copy(dataset.variables['ixseps2'][...]) - jyseps1_1 = np.copy(dataset.variables['jyseps1_1'][...]) - jyseps2_1 = np.copy(dataset.variables['jyseps2_1'][...]) - jyseps1_2 = np.copy(dataset.variables['jyseps1_2'][...]) - jyseps2_2 = np.copy(dataset.variables['jyseps2_2'][...]) - ny_inner = np.copy(dataset.variables['ny_inner'][...]) + y_boundary_guards = np.copy(dataset.variables["y_boundary_guards"][...]) + ixseps1 = np.copy(dataset.variables["ixseps1"][...]) + ixseps2 = np.copy(dataset.variables["ixseps2"][...]) + jyseps1_1 = np.copy(dataset.variables["jyseps1_1"][...]) + jyseps2_1 = np.copy(dataset.variables["jyseps2_1"][...]) + jyseps1_2 = np.copy(dataset.variables["jyseps1_2"][...]) + jyseps2_2 = np.copy(dataset.variables["jyseps2_2"][...]) + ny_inner = np.copy(dataset.variables["ny_inner"][...]) # print("y_boundary_guards",y_boundary_guards) # print("ixseps1",ixseps1) # print("ixseps2",ixseps2) @@ -405,52 +520,61 @@ def plot_corners_get_dmplex_data(file_path,interactive_plot=False,print_cells_to # print("ny_inner",ny_inner) exclude_y_guard_cells = True if exclude_y_guard_cells: - s = slice(0,-1,1), list(range(y_boundary_guards,ny_inner+y_boundary_guards)) + list(range(ny_inner+3*y_boundary_guards,2*ny_inner+3*y_boundary_guards)) + s = ( + slice(0, -1, 1), + list(range(y_boundary_guards, ny_inner + y_boundary_guards)) + + list( + range( + ny_inner + 3 * y_boundary_guards, + 2 * ny_inner + 3 * y_boundary_guards, + ) + ), + ) else: - s = slice(0,-1,1), slice(0,-1,1) - data_Rxy = np.copy(dataset.variables['Rxy_corners'][:])[s] - data_Zxy = np.copy(dataset.variables['Zxy_corners'][:])[s] - data_Rxy_lr = np.copy(dataset.variables['Rxy_lower_right_corners'][:])[s] - data_Zxy_lr = np.copy(dataset.variables['Zxy_lower_right_corners'][:])[s] - data_Rxy_ur = np.copy(dataset.variables['Rxy_upper_right_corners'][:])[s] - data_Zxy_ur = np.copy(dataset.variables['Zxy_upper_right_corners'][:])[s] - data_Rxy_ul = np.copy(dataset.variables['Rxy_upper_left_corners'][:])[s] - data_Zxy_ul = np.copy(dataset.variables['Zxy_upper_left_corners'][:])[s] + s = slice(0, -1, 1), slice(0, -1, 1) + data_Rxy = np.copy(dataset.variables["Rxy_corners"][:])[s] + data_Zxy = np.copy(dataset.variables["Zxy_corners"][:])[s] + data_Rxy_lr = np.copy(dataset.variables["Rxy_lower_right_corners"][:])[s] + data_Zxy_lr = np.copy(dataset.variables["Zxy_lower_right_corners"][:])[s] + data_Rxy_ur = np.copy(dataset.variables["Rxy_upper_right_corners"][:])[s] + data_Zxy_ur = np.copy(dataset.variables["Zxy_upper_right_corners"][:])[s] + data_Rxy_ul = np.copy(dataset.variables["Rxy_upper_left_corners"][:])[s] + data_Zxy_ul = np.copy(dataset.variables["Zxy_upper_left_corners"][:])[s] dataset.close() # check that the points are indeed unique - unique_points_2D(data_Rxy,data_Zxy) - unique_points_2D(data_Rxy_lr,data_Zxy_lr) - unique_points_2D(data_Rxy_ur,data_Zxy_ur) - unique_points_2D(data_Rxy_ul,data_Zxy_ul) + unique_points_2D(data_Rxy, data_Zxy) + unique_points_2D(data_Rxy_lr, data_Zxy_lr) + unique_points_2D(data_Rxy_ur, data_Zxy_ur) + unique_points_2D(data_Rxy_ul, data_Zxy_ul) # construct array with unique set of vertices covering the # Hypnotoad grid. Nx, Ny = data_Rxy.shape - Npoint = Nx*Ny - - corners_Rxy = np.zeros((Nx+1,Ny+2)) - corners_Rxy[0:Nx,0:Ny//2] = data_Rxy[:,0:Ny//2] - corners_Rxy[Nx,0:Ny//2] = data_Rxy_lr[-1,0:Ny//2] - corners_Rxy[1:Nx+1,Ny//2] = data_Rxy_ur[:,Ny//2 - 1] - corners_Rxy[0,Ny//2] = data_Rxy_ul[0,Ny//2 -1] - - corners_Rxy[0:Nx,Ny//2+1:Ny+1] = data_Rxy[:,Ny//2:Ny] - corners_Rxy[Nx,Ny//2+1:Ny+1] = data_Rxy_lr[Nx-1,Ny//2:Ny] - corners_Rxy[0:Nx,Ny+1] = data_Rxy_ul[0:Nx,Ny-1] - corners_Rxy[Nx,Ny+1] = data_Rxy_ur[Nx-1,Ny-1] - - corners_Zxy = np.zeros((Nx+1,Ny+2)) - corners_Zxy[0:Nx,0:Ny//2] = data_Zxy[:,0:Ny//2] - corners_Zxy[Nx,0:Ny//2] = data_Zxy_lr[-1,0:Ny//2] - corners_Zxy[1:Nx+1,Ny//2] = data_Zxy_ur[:,Ny//2 - 1] - corners_Zxy[0,Ny//2] = data_Zxy_ul[0,Ny//2 -1] - - corners_Zxy[0:Nx,Ny//2+1:Ny+1] = data_Zxy[:,Ny//2:Ny] - corners_Zxy[Nx,Ny//2+1:Ny+1] = data_Zxy_lr[Nx-1,Ny//2:Ny] - corners_Zxy[0:Nx,Ny+1] = data_Zxy_ul[0:Nx,Ny-1] - corners_Zxy[Nx,Ny+1] = data_Zxy_ur[Nx-1,Ny-1] + Npoint = Nx * Ny + + corners_Rxy = np.zeros((Nx + 1, Ny + 2)) + corners_Rxy[0:Nx, 0 : Ny // 2] = data_Rxy[:, 0 : Ny // 2] + corners_Rxy[Nx, 0 : Ny // 2] = data_Rxy_lr[-1, 0 : Ny // 2] + corners_Rxy[1 : Nx + 1, Ny // 2] = data_Rxy_ur[:, Ny // 2 - 1] + corners_Rxy[0, Ny // 2] = data_Rxy_ul[0, Ny // 2 - 1] + + corners_Rxy[0:Nx, Ny // 2 + 1 : Ny + 1] = data_Rxy[:, Ny // 2 : Ny] + corners_Rxy[Nx, Ny // 2 + 1 : Ny + 1] = data_Rxy_lr[Nx - 1, Ny // 2 : Ny] + corners_Rxy[0:Nx, Ny + 1] = data_Rxy_ul[0:Nx, Ny - 1] + corners_Rxy[Nx, Ny + 1] = data_Rxy_ur[Nx - 1, Ny - 1] + + corners_Zxy = np.zeros((Nx + 1, Ny + 2)) + corners_Zxy[0:Nx, 0 : Ny // 2] = data_Zxy[:, 0 : Ny // 2] + corners_Zxy[Nx, 0 : Ny // 2] = data_Zxy_lr[-1, 0 : Ny // 2] + corners_Zxy[1 : Nx + 1, Ny // 2] = data_Zxy_ur[:, Ny // 2 - 1] + corners_Zxy[0, Ny // 2] = data_Zxy_ul[0, Ny // 2 - 1] + + corners_Zxy[0:Nx, Ny // 2 + 1 : Ny + 1] = data_Zxy[:, Ny // 2 : Ny] + corners_Zxy[Nx, Ny // 2 + 1 : Ny + 1] = data_Zxy_lr[Nx - 1, Ny // 2 : Ny] + corners_Zxy[0:Nx, Ny + 1] = data_Zxy_ul[0:Nx, Ny - 1] + corners_Zxy[Nx, Ny + 1] = data_Zxy_ur[Nx - 1, Ny - 1] # check that the points are indeed unique - unique_points_2D(corners_Rxy,corners_Zxy) + unique_points_2D(corners_Rxy, corners_Zxy) # Visualise points with scatter plots # Get 1D lists of points for the scatter plot @@ -463,172 +587,268 @@ def plot_corners_get_dmplex_data(file_path,interactive_plot=False,print_cells_to Rpoints_ul = np.reshape(data_Rxy_ul, (Npoint,)) Zpoints_ul = np.reshape(data_Zxy_ul, (Npoint,)) - Npoint_full = (Nx+1)*(Ny+2) + Npoint_full = (Nx + 1) * (Ny + 2) Rpoints_full = np.reshape(corners_Rxy, (Npoint_full,)) Zpoints_full = np.reshape(corners_Zxy, (Npoint_full,)) # remove any duplicate points - Rpoints_full, Zpoints_full = remove_nonunique_points(Rpoints_full,Zpoints_full) - unique_points_1D(Rpoints_full,Zpoints_full) + Rpoints_full, Zpoints_full = remove_nonunique_points(Rpoints_full, Zpoints_full) + unique_points_1D(Rpoints_full, Zpoints_full) - save_ivertex_indices_to_netcdf(file_path, file_path[:-3]+".corners.nc",Rpoints_full,Zpoints_full) + save_ivertex_indices_to_netcdf( + file_path, file_path[:-3] + ".corners.nc", Rpoints_full, Zpoints_full + ) # get an array containing, for each cell, the list of vertices, # labelled by the global index defined implicitly by Rpoints_full, Zpoints_full # store this in `cell_vertices`, and store the RZ values of each vertex # cell_vertices[icell,ivertex] in R, Z = cell_vertices_RZ[icell,ivertex,:] - cell_vertices, cell_vertices_RZ, vertex_list = get_cell_vertex_list(data_Rxy, data_Zxy, - data_Rxy_lr, data_Zxy_lr, data_Rxy_ur, data_Zxy_ur, - data_Rxy_ul, data_Zxy_ul, Rpoints_full, Zpoints_full) + cell_vertices, cell_vertices_RZ, vertex_list = get_cell_vertex_list( + data_Rxy, + data_Zxy, + data_Rxy_lr, + data_Zxy_lr, + data_Rxy_ur, + data_Zxy_ur, + data_Rxy_ul, + data_Zxy_ul, + Rpoints_full, + Zpoints_full, + ) ncells, nvertices = np.shape(cell_vertices) if print_cells_to_screen_output: - for icell in range(0,ncells): - print(cell_vertices[icell,:]) - print(cell_vertices_RZ[icell,:,:]) + for icell in range(0, ncells): + print(cell_vertices[icell, :]) + print(cell_vertices_RZ[icell, :, :]) # test getting data for resaving to hypnotoad netcdf file - ivertex_corners = get_boutxx_corner_index(data_Rxy,data_Zxy,Rpoints_full,Zpoints_full) - ivertex_corners_lr = get_boutxx_corner_index(data_Rxy_lr,data_Zxy_lr,Rpoints_full,Zpoints_full) - ivertex_corners_ul = get_boutxx_corner_index(data_Rxy_ul,data_Zxy_ul,Rpoints_full,Zpoints_full) - ivertex_corners_ur = get_boutxx_corner_index(data_Rxy_ur,data_Zxy_ur,Rpoints_full,Zpoints_full) - #print("ivertex_corners",ivertex_corners) - #print("ivertex_corners_lr",ivertex_corners_lr) - #print("ivertex_corners_ul",ivertex_corners_ul) - #print("ivertex_corners_ur",ivertex_corners_ur) - ivertex_pfr_lower, Rxy_pfr_lower, Zxy_pfr_lower = get_pfr_lower_boundary_vertices(ivertex_corners,data_Rxy,data_Zxy, - ivertex_corners_ul,data_Rxy_ul,data_Zxy_ul, - y_boundary_guards,jyseps1_1,exclude_y_guard_cells=exclude_y_guard_cells) - - ivertex_pfr_upper, Rxy_pfr_upper, Zxy_pfr_upper = get_pfr_upper_boundary_vertices(ivertex_corners,data_Rxy,data_Zxy, - ivertex_corners_ul,data_Rxy_ul,data_Zxy_ul, - y_boundary_guards,jyseps1_2,jyseps2_1,jyseps2_2, - ny_inner, exclude_y_guard_cells=exclude_y_guard_cells) - - ivertex_sol_vac_left, Rxy_sol_vac_left, Zxy_sol_vac_left = get_vac_left_boundary_vertices(ivertex_corners_lr,data_Rxy_lr,data_Zxy_lr, - ivertex_corners_ur,data_Rxy_ur,data_Zxy_ur, - y_boundary_guards,ny_inner) - - ivertex_sol_vac_right, Rxy_sol_vac_right, Zxy_sol_vac_right = get_vac_right_boundary_vertices(ivertex_corners_lr,data_Rxy_lr,data_Zxy_lr, - ivertex_corners_ur,data_Rxy_ur,data_Zxy_ur, - jyseps1_2,ny_inner) - - ivertex_core, Rxy_core, Zxy_core = get_core_boundary_vertices(ivertex_corners, data_Rxy, data_Zxy, - jyseps1_1, jyseps2_1, jyseps1_2, jyseps2_2, - y_boundary_guards, exclude_y_guard_cells=exclude_y_guard_cells) - - ivertex_target_ll, Rxy_target_ll, Zxy_target_ll = get_target_ll_vertices(ivertex_corners, data_Rxy, data_Zxy, - ivertex_corners_lr, data_Rxy_lr, data_Zxy_lr) - - ivertex_target_ul, Rxy_target_ul, Zxy_target_ul = get_target_ul_vertices(ivertex_corners_ul, data_Rxy_ul, data_Zxy_ul, - ivertex_corners_ur, data_Rxy_ur, data_Zxy_ur, - ny_inner) - - ivertex_target_ur, Rxy_target_ur, Zxy_target_ur = get_target_ur_vertices(ivertex_corners, data_Rxy, data_Zxy, - ivertex_corners_lr, data_Rxy_lr, data_Zxy_lr, - ny_inner) - - ivertex_target_lr, Rxy_target_lr, Zxy_target_lr = get_target_lr_vertices(ivertex_corners_ur, data_Rxy_ur, data_Zxy_ur, - ivertex_corners_ul, data_Rxy_ul, data_Zxy_ul) + ivertex_corners = get_boutxx_corner_index( + data_Rxy, data_Zxy, Rpoints_full, Zpoints_full + ) + ivertex_corners_lr = get_boutxx_corner_index( + data_Rxy_lr, data_Zxy_lr, Rpoints_full, Zpoints_full + ) + ivertex_corners_ul = get_boutxx_corner_index( + data_Rxy_ul, data_Zxy_ul, Rpoints_full, Zpoints_full + ) + ivertex_corners_ur = get_boutxx_corner_index( + data_Rxy_ur, data_Zxy_ur, Rpoints_full, Zpoints_full + ) + # print("ivertex_corners",ivertex_corners) + # print("ivertex_corners_lr",ivertex_corners_lr) + # print("ivertex_corners_ul",ivertex_corners_ul) + # print("ivertex_corners_ur",ivertex_corners_ur) + ivertex_pfr_lower, Rxy_pfr_lower, Zxy_pfr_lower = get_pfr_lower_boundary_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + ivertex_corners_ul, + data_Rxy_ul, + data_Zxy_ul, + y_boundary_guards, + jyseps1_1, + exclude_y_guard_cells=exclude_y_guard_cells, + ) + + ivertex_pfr_upper, Rxy_pfr_upper, Zxy_pfr_upper = get_pfr_upper_boundary_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + ivertex_corners_ul, + data_Rxy_ul, + data_Zxy_ul, + y_boundary_guards, + jyseps1_2, + jyseps2_1, + jyseps2_2, + ny_inner, + exclude_y_guard_cells=exclude_y_guard_cells, + ) + + ivertex_sol_vac_left, Rxy_sol_vac_left, Zxy_sol_vac_left = ( + get_vac_left_boundary_vertices( + ivertex_corners_lr, + data_Rxy_lr, + data_Zxy_lr, + ivertex_corners_ur, + data_Rxy_ur, + data_Zxy_ur, + y_boundary_guards, + ny_inner, + ) + ) + + ivertex_sol_vac_right, Rxy_sol_vac_right, Zxy_sol_vac_right = ( + get_vac_right_boundary_vertices( + ivertex_corners_lr, + data_Rxy_lr, + data_Zxy_lr, + ivertex_corners_ur, + data_Rxy_ur, + data_Zxy_ur, + jyseps1_2, + ny_inner, + ) + ) + + ivertex_core, Rxy_core, Zxy_core = get_core_boundary_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + jyseps1_1, + jyseps2_1, + jyseps1_2, + jyseps2_2, + y_boundary_guards, + exclude_y_guard_cells=exclude_y_guard_cells, + ) + + ivertex_target_ll, Rxy_target_ll, Zxy_target_ll = get_target_ll_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + ivertex_corners_lr, + data_Rxy_lr, + data_Zxy_lr, + ) + + ivertex_target_ul, Rxy_target_ul, Zxy_target_ul = get_target_ul_vertices( + ivertex_corners_ul, + data_Rxy_ul, + data_Zxy_ul, + ivertex_corners_ur, + data_Rxy_ur, + data_Zxy_ur, + ny_inner, + ) + + ivertex_target_ur, Rxy_target_ur, Zxy_target_ur = get_target_ur_vertices( + ivertex_corners, + data_Rxy, + data_Zxy, + ivertex_corners_lr, + data_Rxy_lr, + data_Zxy_lr, + ny_inner, + ) + + ivertex_target_lr, Rxy_target_lr, Zxy_target_lr = get_target_lr_vertices( + ivertex_corners_ur, + data_Rxy_ur, + data_Zxy_ur, + ivertex_corners_ul, + data_Rxy_ul, + data_Zxy_ul, + ) boundary_vertex_info = { - "pfr_lower" : {"ivertex": ivertex_pfr_lower, - "Rxy" : Rxy_pfr_lower, - "Zxy" : Zxy_pfr_lower, - "color" : "g", - "marker" : "1", - "DMFaceSetsLabel" : 100, - }, - "pfr_upper" : {"ivertex": ivertex_pfr_upper, - "Rxy" : Rxy_pfr_upper, - "Zxy" : Zxy_pfr_upper, - "color" : "b", - "marker" : "2", - "DMFaceSetsLabel" : 100, - }, - "sol_vac_left" : {"ivertex": ivertex_sol_vac_left, - "Rxy" : Rxy_sol_vac_left, - "Zxy" : Zxy_sol_vac_left, - "color" : "r", - "marker" : "3", - "DMFaceSetsLabel" : 200, - }, - "sol_vac_right" : {"ivertex": ivertex_sol_vac_right, - "Rxy" : Rxy_sol_vac_right, - "Zxy" : Zxy_sol_vac_right, - "color" : "r", - "marker" : "3", - "DMFaceSetsLabel" : 200, - }, - "core" : {"ivertex": ivertex_core, - "Rxy" : Rxy_core, - "Zxy" : Zxy_core, - "color" : "k", - "marker" : "4", - "DMFaceSetsLabel" : 300, - }, - "target_ll" : {"ivertex": ivertex_target_ll, - "Rxy" : Rxy_target_ll, - "Zxy" : Zxy_target_ll, - "color" : "y", - "marker" : "3", - "DMFaceSetsLabel" : 400, - }, - "target_ul" : {"ivertex": ivertex_target_ul, - "Rxy" : Rxy_target_ul, - "Zxy" : Zxy_target_ul, - "color" : "y", - "marker" : "3", - "DMFaceSetsLabel" : 400, - }, - "target_ur" : {"ivertex": ivertex_target_ur, - "Rxy" : Rxy_target_ur, - "Zxy" : Zxy_target_ur, - "color" : "y", - "marker" : "3", - "DMFaceSetsLabel" : 400, - }, - "target_lr" : {"ivertex": ivertex_target_lr, - "Rxy" : Rxy_target_lr, - "Zxy" : Zxy_target_lr, - "color" : "y", - "marker" : "3", - "DMFaceSetsLabel" : 400, - }, + "pfr_lower": { + "ivertex": ivertex_pfr_lower, + "Rxy": Rxy_pfr_lower, + "Zxy": Zxy_pfr_lower, + "color": "g", + "marker": "1", + "DMFaceSetsLabel": 100, + }, + "pfr_upper": { + "ivertex": ivertex_pfr_upper, + "Rxy": Rxy_pfr_upper, + "Zxy": Zxy_pfr_upper, + "color": "b", + "marker": "2", + "DMFaceSetsLabel": 100, + }, + "sol_vac_left": { + "ivertex": ivertex_sol_vac_left, + "Rxy": Rxy_sol_vac_left, + "Zxy": Zxy_sol_vac_left, + "color": "r", + "marker": "3", + "DMFaceSetsLabel": 200, + }, + "sol_vac_right": { + "ivertex": ivertex_sol_vac_right, + "Rxy": Rxy_sol_vac_right, + "Zxy": Zxy_sol_vac_right, + "color": "r", + "marker": "3", + "DMFaceSetsLabel": 200, + }, + "core": { + "ivertex": ivertex_core, + "Rxy": Rxy_core, + "Zxy": Zxy_core, + "color": "k", + "marker": "4", + "DMFaceSetsLabel": 300, + }, + "target_ll": { + "ivertex": ivertex_target_ll, + "Rxy": Rxy_target_ll, + "Zxy": Zxy_target_ll, + "color": "y", + "marker": "3", + "DMFaceSetsLabel": 400, + }, + "target_ul": { + "ivertex": ivertex_target_ul, + "Rxy": Rxy_target_ul, + "Zxy": Zxy_target_ul, + "color": "y", + "marker": "3", + "DMFaceSetsLabel": 400, + }, + "target_ur": { + "ivertex": ivertex_target_ur, + "Rxy": Rxy_target_ur, + "Zxy": Zxy_target_ur, + "color": "y", + "marker": "3", + "DMFaceSetsLabel": 400, + }, + "target_lr": { + "ivertex": ivertex_target_lr, + "Rxy": Rxy_target_lr, + "Zxy": Zxy_target_lr, + "color": "y", + "marker": "3", + "DMFaceSetsLabel": 400, + }, } # Make a scatter plot to show the mesh corners plt.figure(figsize=(10, 6)) x = Rpoints y = Zpoints - #scatter = plt.scatter(x, y, c='b',marker='1') - #scatter = plt.scatter(Rpoints_lr, Zpoints_lr, c='r',marker='2') - #scatter = plt.scatter(Rpoints_ur, Zpoints_ur, c='g',marker='3') - #scatter = plt.scatter(Rpoints_ul, Zpoints_ul, c='k',marker='4') - scatter = plt.scatter(Rpoints_full, Zpoints_full, c='m',marker='x') + # scatter = plt.scatter(x, y, c='b',marker='1') + # scatter = plt.scatter(Rpoints_lr, Zpoints_lr, c='r',marker='2') + # scatter = plt.scatter(Rpoints_ur, Zpoints_ur, c='g',marker='3') + # scatter = plt.scatter(Rpoints_ul, Zpoints_ul, c='k',marker='4') + scatter = plt.scatter(Rpoints_full, Zpoints_full, c="m", marker="x") for key in boundary_vertex_info.keys(): Rxy = boundary_vertex_info[key]["Rxy"] Zxy = boundary_vertex_info[key]["Zxy"] color = boundary_vertex_info[key]["color"] marker = boundary_vertex_info[key]["marker"] - scatter = plt.scatter(Rxy,Zxy, c=color,marker=marker) + scatter = plt.scatter(Rxy, Zxy, c=color, marker=marker) # uncomment for labels on original data points - #for ic in range(0,Npoint): + # for ic in range(0,Npoint): # plt.text(x[ic],y[ic],str(ic)) # plt.text(Rpoints_ul[ic],Zpoints_ul[ic],str(ic)) # uncomment for labels on aggregated array of points - #Npoint_full = len(Rpoints_full) - #for ic in range(0,Npoint_full): + # Npoint_full = len(Rpoints_full) + # for ic in range(0,Npoint_full): # plt.text(Rpoints_full[ic],Zpoints_full[ic],str(ic)) for key in boundary_vertex_info.keys(): Rxy = boundary_vertex_info[key]["Rxy"] Zxy = boundary_vertex_info[key]["Zxy"] ivertex = boundary_vertex_info[key]["ivertex"] Nvertex = len(ivertex) - for i in range(0,Nvertex): - plt.text(Rxy[i],Zxy[i],str(ivertex[i])) + for i in range(0, Nvertex): + plt.text(Rxy[i], Zxy[i], str(ivertex[i])) - plt.title('Meshpoints') - plt.xlabel('R') - plt.ylabel('Z') + plt.title("Meshpoints") + plt.xlabel("R") + plt.ylabel("Z") if interactive_plot: plt.show() - plt.savefig(file_path+".mesh_plot.pdf") - return Nx, Ny, cell_vertices, vertex_list, boundary_vertex_info \ No newline at end of file + plt.savefig(file_path + ".mesh_plot.pdf") + return Nx, Ny, cell_vertices, vertex_list, boundary_vertex_info diff --git a/tests/integrated/dmplex-vertex-coordinates/runtest b/tests/integrated/dmplex-vertex-coordinates/runtest index f46df2546..980655a09 100755 --- a/tests/integrated/dmplex-vertex-coordinates/runtest +++ b/tests/integrated/dmplex-vertex-coordinates/runtest @@ -41,7 +41,7 @@ grid_filenames = [ "example_udn.grd.nc", "example_udn2.grd.nc", "example_usn.grd.nc", - ] +] ### Define functions @@ -63,17 +63,13 @@ def get_DMPlex_vertices(dmplex_path): ## Get vertices # vertices are the simplest object, so depth = 0 - vertex_start, vertex_end = dm.getDepthStratum( - 0 - ) + vertex_start, vertex_end = dm.getDepthStratum(0) R = [] Z = [] # Find offset to locate matching coordinates to point p for p in range(vertex_start, vertex_end): - off = section.getOffset( - p - ) + off = section.getOffset(p) R.append(coords[off]) Z.append(coords[off + 1]) @@ -104,7 +100,6 @@ def get_hypnotoad_vertices(gridfile): ### Prepare for test for gridfile in grid_filenames: - if verbose: print("\n", "*" * 30, "\n") print(f"Testing {gridfile}") @@ -135,10 +130,9 @@ for gridfile in grid_filenames: tmp_path = gridzipfile_path.with_name(gridzipfile_path.name + ".tmp") with urllib.request.urlopen(gridurl, timeout=60) as response: - if response.status != 200: raise RuntimeError(f"Test grid download failed: HTTP {response.status}") - + # Copy bits of the file from response to a temp file # This ensures no partial files are left if the download fails with open(tmp_path, "wb") as out_file: @@ -148,7 +142,6 @@ for gridfile in grid_filenames: tmp_path.replace(gridzipfile_path) with zipfile.ZipFile(gridzipfile_path, "r") as zf: - zip_contents = set(zf.namelist()) try: # Extract only expected grids @@ -164,7 +157,9 @@ for gridfile in grid_filenames: file_hash = hashlib.sha256(f.read()).hexdigest() if file_hash != expected_hash: - raise RuntimeError("Downloaded grid zip file hash does not match expected value") + raise RuntimeError( + "Downloaded grid zip file hash does not match expected value" + ) ## Run Hermes-3 to generate DMPlex mesh # Find how many cores to use @@ -205,7 +200,7 @@ for gridfile in grid_filenames: # Extract dmplex_vertices = get_DMPlex_vertices(dmplex_path) hypno_vertices = get_hypnotoad_vertices(gridfile) - + R_dmplex = dmplex_vertices[:, 0] Z_dmplex = dmplex_vertices[:, 1] R_hypno = hypno_vertices[:, 0] @@ -218,19 +213,23 @@ for gridfile in grid_filenames: # Hypnotoad has duplicate vertices due to overlapping vertices between cells. # The vertices are not in an exactly identical position due to finite precision # on the equilibrium interpolation done in Hypnotoad. - # The test just ensures that every vertex in one mesh has a matching vertex in + # The test just ensures that every vertex in one mesh has a matching vertex in # the other mesh within the tolerance. - + # Ensure all Hypnotoad vertices exist in DMPlex for i in range(len(R_hypno)): - nearest_distance = np.min(np.sqrt((R_dmplex - R_hypno[i])**2 + (Z_dmplex - Z_hypno[i])**2)) - + nearest_distance = np.min( + np.sqrt((R_dmplex - R_hypno[i]) ** 2 + (Z_dmplex - Z_hypno[i]) ** 2) + ) + if not nearest_distance < tolerance: raise AssertionError("Found Hypnotoad vertex with no match in DMPlex") - + # Ensure all DMPlex vertices exist in Hypnotoad for i in range(len(R_dmplex)): - nearest_distance = np.min(np.sqrt((R_hypno - R_dmplex[i])**2 + (Z_hypno - Z_dmplex[i])**2)) - + nearest_distance = np.min( + np.sqrt((R_hypno - R_dmplex[i]) ** 2 + (Z_hypno - Z_dmplex[i]) ** 2) + ) + if not nearest_distance < tolerance: raise AssertionError("Found DMPlex vertex with no match in Hypnotoad") diff --git a/tests/integrated/particle-pusher/runtest b/tests/integrated/particle-pusher/runtest index b04b96a2d..417f2c012 100755 --- a/tests/integrated/particle-pusher/runtest +++ b/tests/integrated/particle-pusher/runtest @@ -16,10 +16,19 @@ shell("ln -s ../../../bout-particle-push bout-particle-push") # make the run directory shell("mkdir particle-push-slab-test") + # A function to define the BOUT.inp file contents -def particle_push_input(nx=40, ny=36, dt=0.005, nsteps=1, - iz_rate=0.0, rec_rate=0.0, remove_threshold=0.0, merge_threshold=0.0): - input_file_string = f""" +def particle_push_input( + nx=40, + ny=36, + dt=0.005, + nsteps=1, + iz_rate=0.0, + rec_rate=0.0, + remove_threshold=0.0, + merge_threshold=0.0, +): + input_file_string = f""" # run a simulation with this file using # $ OMP_NUM_THREADS=1 mpirun -np 8 ./bout-particle-push -d particle-push-slab # postprocess with @@ -71,114 +80,174 @@ def particle_push_input(nx=40, ny=36, dt=0.005, nsteps=1, iz_rate = {iz_rate} rec_rate = {rec_rate} """ - return input_file_string + return input_file_string + # function for reading resulting particle data def load_particle_data(file_path): - particle_data = h5py.File(file_path,"r") - data_all_timesteps = list(particle_data.keys()) - nstep = len(data_all_timesteps) - - particle_positions = [] - for it in range(0,nstep): - try: - group = particle_data[f"Step#{it}"] - #print(list(group.keys())) - P_0 = group["POSITION_0"] - P_1 = group["POSITION_1"] - nparticles = len(P_0) - pdata = np.zeros((nparticles,2)) - pdata[:,0] = P_0 - pdata[:,1] = P_1 - particle_positions.append(pdata) - except (KeyError) as error: - print(f"No particles at time step {it}: {error}") - # assign empty particle data - pdata = np.empty((0,2)) - particle_positions.append(pdata) - return particle_positions - -def test_particle_push(file, mode, nsteps = 10, remove_threshold = 0.0, merge_threshold = 0.0, iz_rate = 0.0, rec_rate = 0.0): - """ - Set up particle push test and pass settings to the input file. - - Test mode can be: - - advection: ionisation and recombination rates are 0. Particle count is checked to be conserved. - - ionisation: ionisation rate > 0, recombination rate = 0. Particle count should increase - - recombination: ionisation rate = 0, recombination rate > 0. Particle count should decrease - """ - - with open(file,"w") as file: - file.write(particle_push_input(nx=40, ny=36, dt=0.005, nsteps=nsteps, - iz_rate=iz_rate, rec_rate=rec_rate, remove_threshold=remove_threshold, merge_threshold=merge_threshold)) - - # Command to run - cmd = "./bout-particle-push -d particle-push-slab-test" - nproc = 1 - # Launch using MPI - s, out = launch_safe(cmd, nproc=nproc, mthread=1, pipe=True) - particle_positions_file_path = r"particle-push-slab-test/particle_trajectories.h5part" - particle_positions = load_particle_data(particle_positions_file_path) - nsteps_out = len(particle_positions) - - # number of steps the same - np.testing.assert_equal( - nsteps, nsteps_out, - err_msg=f"Number of steps in outputdoes not match expected for {mode} test") - - # number of markers (or super particles) conserved - nparticles_0, _ = np.shape(particle_positions[0]) - nparticles_N, _ = np.shape(particle_positions[-1]) - if mode == "advection": - np.testing.assert_equal( - nparticles_0, nparticles_N, - err_msg="Number of particles not conserved in advection test") - - - # check the mass diagnostic - ncfile_path = r"particle-push-slab-test/BOUT.dmp.0.nc" - - with nc.Dataset(ncfile_path, mode='r') as ncdataset: - total_mass = np.copy(ncdataset.variables["total_mass"][:]) - total_ion_mass = np.copy(ncdataset.variables["total_ion_mass"][:]) - total_neutral_mass = np.copy(ncdataset.variables["total_neutral_mass"][:]) - - # total mass should always be conserved - np.testing.assert_allclose(total_mass[0], total_mass[-1], atol = 1.0e-11, - err_msg=f"Total mass not conserved in {mode} test") - - # if we ionise particles there should be mass exchange between neutral and ion - if mode == "ionisation": - np.testing.assert_array_less( - total_ion_mass[0], total_ion_mass[-1], - err_msg="Ion mass did not increase in ionisation test") - np.testing.assert_array_less( - total_neutral_mass[-1], total_neutral_mass[0], - err_msg="Neutral mass did not decrease in ionisation test") - - # Recombination is the opposite - if mode == "recombination": - np.testing.assert_array_less( - total_ion_mass[-1], total_ion_mass[0], - err_msg="Ion mass did not decrease in recombination test") - np.testing.assert_array_less( - total_neutral_mass[0], total_neutral_mass[-1], - err_msg="Neutral mass did not increase in recombination test") + particle_data = h5py.File(file_path, "r") + data_all_timesteps = list(particle_data.keys()) + nstep = len(data_all_timesteps) + + particle_positions = [] + for it in range(0, nstep): + try: + group = particle_data[f"Step#{it}"] + # print(list(group.keys())) + P_0 = group["POSITION_0"] + P_1 = group["POSITION_1"] + nparticles = len(P_0) + pdata = np.zeros((nparticles, 2)) + pdata[:, 0] = P_0 + pdata[:, 1] = P_1 + particle_positions.append(pdata) + except KeyError as error: + print(f"No particles at time step {it}: {error}") + # assign empty particle data + pdata = np.empty((0, 2)) + particle_positions.append(pdata) + return particle_positions + + +def test_particle_push( + file, + mode, + nsteps=10, + remove_threshold=0.0, + merge_threshold=0.0, + iz_rate=0.0, + rec_rate=0.0, +): + """ + Set up particle push test and pass settings to the input file. + + Test mode can be: + - advection: ionisation and recombination rates are 0. Particle count is checked to be conserved. + - ionisation: ionisation rate > 0, recombination rate = 0. Particle count should increase + - recombination: ionisation rate = 0, recombination rate > 0. Particle count should decrease + """ + + with open(file, "w") as file: + file.write( + particle_push_input( + nx=40, + ny=36, + dt=0.005, + nsteps=nsteps, + iz_rate=iz_rate, + rec_rate=rec_rate, + remove_threshold=remove_threshold, + merge_threshold=merge_threshold, + ) + ) + + # Command to run + cmd = "./bout-particle-push -d particle-push-slab-test" + nproc = 1 + # Launch using MPI + s, out = launch_safe(cmd, nproc=nproc, mthread=1, pipe=True) + particle_positions_file_path = ( + r"particle-push-slab-test/particle_trajectories.h5part" + ) + particle_positions = load_particle_data(particle_positions_file_path) + nsteps_out = len(particle_positions) + + # number of steps the same + np.testing.assert_equal( + nsteps, + nsteps_out, + err_msg=f"Number of steps in outputdoes not match expected for {mode} test", + ) + + # number of markers (or super particles) conserved + nparticles_0, _ = np.shape(particle_positions[0]) + nparticles_N, _ = np.shape(particle_positions[-1]) + if mode == "advection": + np.testing.assert_equal( + nparticles_0, + nparticles_N, + err_msg="Number of particles not conserved in advection test", + ) + + # check the mass diagnostic + ncfile_path = r"particle-push-slab-test/BOUT.dmp.0.nc" + + with nc.Dataset(ncfile_path, mode="r") as ncdataset: + total_mass = np.copy(ncdataset.variables["total_mass"][:]) + total_ion_mass = np.copy(ncdataset.variables["total_ion_mass"][:]) + total_neutral_mass = np.copy(ncdataset.variables["total_neutral_mass"][:]) + + # total mass should always be conserved + np.testing.assert_allclose( + total_mass[0], + total_mass[-1], + atol=1.0e-11, + err_msg=f"Total mass not conserved in {mode} test", + ) + + # if we ionise particles there should be mass exchange between neutral and ion + if mode == "ionisation": + np.testing.assert_array_less( + total_ion_mass[0], + total_ion_mass[-1], + err_msg="Ion mass did not increase in ionisation test", + ) + np.testing.assert_array_less( + total_neutral_mass[-1], + total_neutral_mass[0], + err_msg="Neutral mass did not decrease in ionisation test", + ) + + # Recombination is the opposite + if mode == "recombination": + np.testing.assert_array_less( + total_ion_mass[-1], + total_ion_mass[0], + err_msg="Ion mass did not decrease in recombination test", + ) + np.testing.assert_array_less( + total_neutral_mass[0], + total_neutral_mass[-1], + err_msg="Neutral mass did not increase in recombination test", + ) file = "particle-push-slab-test/BOUT.inp" -test_particle_push(file, "advection", nsteps = 3, remove_threshold = 1.0e-10, merge_threshold = 0, iz_rate = 0.00, rec_rate = 0.00) +test_particle_push( + file, + "advection", + nsteps=3, + remove_threshold=1.0e-10, + merge_threshold=0, + iz_rate=0.00, + rec_rate=0.00, +) if verbose: - print("Advection test passed") - -test_particle_push(file, "ionisation", nsteps = 3, remove_threshold = 1.0e-10, merge_threshold = 0, iz_rate = 0.01, rec_rate = 0.00) + print("Advection test passed") + +test_particle_push( + file, + "ionisation", + nsteps=3, + remove_threshold=1.0e-10, + merge_threshold=0, + iz_rate=0.01, + rec_rate=0.00, +) if verbose: - print("Ionisation test passed") - -test_particle_push(file, "recombination", nsteps = 3, remove_threshold = 1.0e-10, merge_threshold = 0, iz_rate = 0.00, rec_rate = 0.0005) + print("Ionisation test passed") + +test_particle_push( + file, + "recombination", + nsteps=3, + remove_threshold=1.0e-10, + merge_threshold=0, + iz_rate=0.00, + rec_rate=0.0005, +) if verbose: - print("Recombination test passed") + print("Recombination test passed") print(" => Test passed") -