From 7d25d8b058ab48d4a37092d53ca9c7b46f14d3bf Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 18 Nov 2022 12:25:54 +0000 Subject: [PATCH 01/34] Write str as byte-strings in to_restart() Need to write char-arrays rather than 'strings' so that BOUT++ can read them. --- xbout/utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/xbout/utils.py b/xbout/utils.py index 7d8757af..9ea1cbda 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -488,7 +488,13 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over restart_ds[v] = data_variable for v in ds.metadata: if v not in restart_exclude_metadata_vars: - restart_ds[v] = ds.metadata[v] + value = ds.metadata[v] + + if isinstance(value, str): + # Write strings as byte-strings so BOUT++ can read them + value = value.encode() + + restart_ds[v] = value # These variables need to be altered, because they depend on the number of # files and/or the rank of this file. From 56e6bfc02527a931ca1d488eb50e9e474897bd83 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 18 Nov 2022 09:16:11 +0000 Subject: [PATCH 02/34] Also blacklist xarray-2022.11.0 Performance issues do not seem to be fixed yet. --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index d326c1c0..488a934e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -28,7 +28,7 @@ setup_requires = setuptools_scm[toml]>=3.4 setuptools_scm_git_archive install_requires = - xarray>=0.18.0,!=2022.9.0,!=2022.10.0 + xarray>=0.18.0,!=2022.9.0,!=2022.10.0,!=2022.11.0 boutdata>=0.1.4 dask[array]>=2.10.0 gelidum>=0.5.3 From c144d5715e975ce677704c9c108fcac0e1a55dfd Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 2 Dec 2022 18:47:08 +0000 Subject: [PATCH 03/34] Allow opening data from BOUT-3.x simulations Some workarounds for variables that were not saved by BOUT++ v3.x. --- xbout/load.py | 70 ++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 55 insertions(+), 15 deletions(-) diff --git a/xbout/load.py b/xbout/load.py index 0c2ef698..fc93a603 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -600,17 +600,40 @@ def _auto_open_mfboutdataset( paths_grid, concat_dims = _arrange_for_concatenation(filepaths, nxpe, nype) - ds = xr.open_mfdataset( - paths_grid, - concat_dim=concat_dims, - combine="nested", - data_vars=data_vars, - preprocess=_preprocess, - engine=filetype, - chunks=chunks, - join="exact", - **kwargs, - ) + try: + ds = xr.open_mfdataset( + paths_grid, + concat_dim=concat_dims, + combine="nested", + data_vars=data_vars, + preprocess=_preprocess, + engine=filetype, + chunks=chunks, + join="exact", + **kwargs, + ) + except ValueError as e: + message_to_catch = ( + "some variables in data_vars are not data variables on the first " + "dataset:" + ) + if str(e)[: len(message_to_catch)] == message_to_catch: + # Open concatenating any variables that are different in + # different files as a work around to support opening older + # data. + ds = xr.open_mfdataset( + paths_grid, + concat_dim=concat_dims, + combine="nested", + data_vars="different", + preprocess=_preprocess, + engine=filetype, + chunks=chunks, + join="exact", + **kwargs, + ) + else: + raise else: # datapath was nested list of Datasets @@ -744,8 +767,16 @@ def get_nonnegative_scalar(ds, key, default=1, info=True): # Check whether this is a single file squashed from the multiple output files of a # parallel run (i.e. NXPE*NYPE > 1 even though there is only a single file to read). - nx = ds["nx"].values - ny = ds["ny"].values + if "nx" in ds: + nx = ds["nx"].values + else: + # Workaround for older data files + nx = ds["MXSUB"].values * ds["NXPE"].values + 2 * ds["MXG"].values + if "ny" in ds: + ny = ds["ny"].values + else: + # Workaround for older data files + ny = ds["MYSUB"].values * ds["NYPE"].values nx_file = ds.dims["x"] ny_file = ds.dims["y"] is_squashed_doublenull = False @@ -758,7 +789,10 @@ def get_nonnegative_scalar(ds, key, default=1, info=True): mxg = 0 # Check if there are two divertor targets present - if ds["jyseps1_2"] > ds["jyseps2_1"]: + # Note: if jyseps2_1 and jyseps1_2 are not in ds it probably + # indicates older data and likely the upper target boundary cells + # were not saved anyway, so continue as if they were not. + if "jyseps2_1" in ds and ds["jyseps1_2"] > ds["jyseps2_1"]: upper_target_cells = myg else: upper_target_cells = 0 @@ -771,7 +805,13 @@ def get_nonnegative_scalar(ds, key, default=1, info=True): nxpe = 1 nype = 1 - is_squashed_doublenull = (ds["jyseps2_1"] != ds["jyseps1_2"]).values + if "jyseps2_1" in ds: + is_squashed_doublenull = (ds["jyseps2_1"] != ds["jyseps1_2"]).values + else: + # For older data with no jyseps2_1 or jyseps1_2 in the + # dataset, probably do not need to handle double null data + # squashed with upper target points. + is_squashed_doublenull = False elif ny_file == ny + 2 * myg: # Older squashed file from double-null grid but containing only lower # target boundary cells. From 81285becbedc68cd0687544d563bfe76229971c5 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 2 Dec 2022 19:22:51 +0000 Subject: [PATCH 04/34] Workarounds to allow xbout methods to work with BOUT-3.x data --- xbout/load.py | 49 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/xbout/load.py b/xbout/load.py index fc93a603..4b8853ce 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -35,6 +35,17 @@ "MYPE", ] _BOUT_TIME_DEPENDENT_META_VARS = ["iteration", "hist_hi", "tt"] +_BOUT_GEOMETRY_VARS = [ + "ixseps1", + "ixseps2", + "jyseps1_1", + "jyseps2_1", + "jyseps1_2", + "jyseps2_2", + "nx", + "ny", + "ny_inner", +] # This code should run whenever any function from this module is imported @@ -350,6 +361,10 @@ def attrs_remove_section(obj, section): pass else: raise ValueError(msg) + for v in _BOUT_GEOMETRY_VARS: + if v not in ds.metadata and v in grid: + ds.metadata[v] = grid[v].values + # Update coordinates to match particular geometry of grid ds = geometries.apply_geometry(ds, geometry, grid=grid) @@ -365,6 +380,40 @@ def attrs_remove_section(obj, section): # BOUT++ ds.bout.fine_interpolation_factor = 8 + if ds.metadata["BOUT_VERSION"] < 4.0: + # Add workarounds for missing information or different conventions in data saved + # by BOUT++ v3.x. + for v in ds: + if ds.metadata["bout_zdim"] in ds[v].dims: + # All fields saved on aligned grid for BOUT-3 + ds[v].attrs["direction_y"] = "Aligned" + + added_location = False + if any( + d in ds[v].dims + for d in ( + ds.metadata["bout_xdim"], + ds.metadata["bout_ydim"], + ds.metadata["bout_zdim"], + ) + ): + # zShift, etc. did not support staggered grids in BOUT++ v3 anyway, so + # just treat all variables as if they were at CELL_CENTRE + ds[v].attrs["cell_location"] = "CELL_CENTRE" + added_location = True + if added_location: + warn( + "Detected data from BOUT++ v3.x. Treating all variables as being " + "at `CELL_CENTRE`. Should be similar to what BOUT++ v3.x did, but " + "if your code uses staggered grids, this may produce unexpected " + "effects in some places." + ) + + if "nz" not in ds.metadata: + # `nz` used to be stored as `MZ` and `MZ` used to include an extra buffer + # point that was not used for data. + ds.metadata["nz"] = ds.metadata["MZ"] - 1 + if info == "terse": print("Read in dataset from {}".format(str(Path(datapath)))) elif info: From 90e532d53b4962a0bd9756cc07854211868adad5 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 4 Dec 2022 16:51:55 +0000 Subject: [PATCH 05/34] Only check BOUT_VERSION for dump or restart files BOUT_VERSION does not have to be defined for grid files (as these are produced by hypnotoad, not BOUT++). --- xbout/load.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/xbout/load.py b/xbout/load.py index 4b8853ce..7c08f9d8 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -380,7 +380,9 @@ def attrs_remove_section(obj, section): # BOUT++ ds.bout.fine_interpolation_factor = 8 - if ds.metadata["BOUT_VERSION"] < 4.0: + if ("dump" in input_type or "restart" in input_type) and ds.metadata[ + "BOUT_VERSION" + ] < 4.0: # Add workarounds for missing information or different conventions in data saved # by BOUT++ v3.x. for v in ds: From 7e4db6b8794cf3db6dda1fbb8bad8c57335a9fe4 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Dec 2022 08:44:52 +0000 Subject: [PATCH 06/34] Blacklist xarray-2022.12.0 Performance issues still not fixed. --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index 488a934e..1f6349b5 100644 --- a/setup.cfg +++ b/setup.cfg @@ -28,7 +28,7 @@ setup_requires = setuptools_scm[toml]>=3.4 setuptools_scm_git_archive install_requires = - xarray>=0.18.0,!=2022.9.0,!=2022.10.0,!=2022.11.0 + xarray>=0.18.0,!=2022.9.0,!=2022.10.0,!=2022.11.0,!=2022.12.0 boutdata>=0.1.4 dask[array]>=2.10.0 gelidum>=0.5.3 From a5046b8d54bc01f25d64402cc67df7df4c138ea1 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Dec 2022 10:29:34 +0000 Subject: [PATCH 07/34] Remove Python-3.7 from the 'master' CI jobs --- .github/workflows/master.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml index c1fc91f2..df706017 100644 --- a/.github/workflows/master.yml +++ b/.github/workflows/master.yml @@ -20,7 +20,7 @@ jobs: if: always() strategy: matrix: - python-version: [3.7, 3.8, 3.9] + python-version: [3.8, 3.9, 3.10] pip-packages: - "setuptools pip pytest pytest-cov coverage codecov boutdata xarray numpy>=1.16.0" fail-fast: false @@ -53,7 +53,7 @@ jobs: if: always() strategy: matrix: - python-version: [3.7, 3.8] + python-version: [3.8] pip-packages: - "setuptools pip pytest pytest-cov coverage codecov boutdata==0.1.4 xarray==0.18.0 dask==2.10.0 numpy==1.18.0 natsort==5.5.0 matplotlib==3.1.1 animatplot==0.4.2 netcdf4==1.4.2 Pillow==6.1.0" # test with oldest supported version of packages. Note, using numpy==1.18.0 as a workaround because numpy==1.17.0 is not supported on Python-3.7, even though we should currently support numpy==1.17.0. fail-fast: false From 207b6ac79d6da7c1dd8fcbc4ce1d85a30d391ae4 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Dec 2022 11:15:09 +0000 Subject: [PATCH 08/34] Add quotes around '3.10' in 'master' CI --- .github/workflows/master.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml index df706017..159d0fb0 100644 --- a/.github/workflows/master.yml +++ b/.github/workflows/master.yml @@ -20,7 +20,7 @@ jobs: if: always() strategy: matrix: - python-version: [3.8, 3.9, 3.10] + python-version: [3.8, 3.9, '3.10'] pip-packages: - "setuptools pip pytest pytest-cov coverage codecov boutdata xarray numpy>=1.16.0" fail-fast: false From fef9ee97972efa11e03e2a00234ff54fc9e571ce Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Dec 2022 13:29:26 +0000 Subject: [PATCH 09/34] Fix bug with outer-x logic in region.get_slices() When adding guard cells, typo meant `xi` was modified instead of `xo`. --- xbout/region.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/region.py b/xbout/region.py index 56634f3e..54e9469c 100644 --- a/xbout/region.py +++ b/xbout/region.py @@ -163,7 +163,7 @@ def get_slices(self, mxg=0, myg=0): xo = self.xouter_ind if self.connection_outer_x is not None and xo is not None: - xi += mxg + xo += mxg yl = self.ylower_ind if self.connection_lower_y is not None and yl is not None: From 7c23536fccc62830390fb2facd0d767f8521b9e2 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 5 Dec 2022 13:33:34 +0000 Subject: [PATCH 10/34] Don't use 'in-place' modification operators in region.py Using `+=` and `-=` (sometimes) modifies input values, which was not intended. Fix is to replace, e.g., `self.nx -= xbndry` with `self.nx = self.nx - xbndry`. --- xbout/region.py | 58 ++++++++++++++++++++++++------------------------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/xbout/region.py b/xbout/region.py index 54e9469c..17bf6454 100644 --- a/xbout/region.py +++ b/xbout/region.py @@ -72,30 +72,30 @@ def __init__( if ds.metadata["keep_xboundaries"]: xbndry = ds.metadata["MXG"] if self.connection_inner_x is None: - self.nx -= xbndry + self.nx = self.nx - xbndry # used to calculate x-coordinate of inner side (self.xinner) - xinner_ind += xbndry + xinner_ind = xinner_ind + xbndry if self.connection_outer_x is None: - self.nx -= xbndry + self.nx = self.nx - xbndry # used to calculate x-coordinate of outer side (self.xouter) - xouter_ind -= xbndry + xouter_ind = xouter_ind - xbndry if ds.metadata["keep_yboundaries"]: ybndry = ds.metadata["MYG"] if self.connection_lower_y is None: - self.ny -= ybndry + self.ny = self.ny - ybndry # used to calculate y-coordinate of lower side (self.ylower) - ylower_ind += ybndry + ylower_ind = ylower_ind + ybndry if self.connection_upper_y is None: - self.ny -= ybndry + self.ny = self.ny - ybndry # used to calculate y-coordinate of upper side (self.yupper) - yupper_ind -= ybndry + yupper_ind = yupper_ind - ybndry # calculate start and end coordinates ##################################### @@ -159,19 +159,19 @@ def get_slices(self, mxg=0, myg=0): """ xi = self.xinner_ind if self.connection_inner_x is not None and xi is not None: - xi -= mxg + xi = xi - mxg xo = self.xouter_ind if self.connection_outer_x is not None and xo is not None: - xo += mxg + xo = xo + mxg yl = self.ylower_ind if self.connection_lower_y is not None and yl is not None: - yl -= myg + yl = yl - myg yu = self.yupper_ind if self.connection_upper_y is not None and yu is not None: - yu += myg + yu = yu + myg return {self.xcoord: slice(xi, xo), self.ycoord: slice(yl, yu)} @@ -189,10 +189,10 @@ def get_inner_guards_slices(self, *, mxg, myg=0): """ ylower = self.ylower_ind if self.connection_lower_y is not None: - ylower -= myg + ylower = ylower - myg yupper = self.yupper_ind if self.connection_upper_y is not None: - yupper += myg + yupper = yupper + myg return { self.xcoord: slice(self.xinner_ind - mxg, self.xinner_ind), self.ycoord: slice(ylower, yupper), @@ -212,10 +212,10 @@ def get_outer_guards_slices(self, *, mxg, myg=0): """ ylower = self.ylower_ind if self.connection_lower_y is not None: - ylower -= myg + ylower = ylower - myg yupper = self.yupper_ind if self.connection_upper_y is not None: - yupper += myg + yupper = yupper + myg return { self.xcoord: slice(self.xouter_ind, self.xouter_ind + mxg), self.ycoord: slice(ylower, yupper), @@ -235,10 +235,10 @@ def get_lower_guards_slices(self, *, myg, mxg=0): """ xinner = self.xinner_ind if self.connection_inner_x is not None: - xinner -= mxg + xinner = xinner - mxg xouter = self.xouter_ind if self.connection_outer_x is not None: - xouter += mxg + xouter = xouter + mxg ystart = self.ylower_ind - myg ystop = self.ylower_ind @@ -270,10 +270,10 @@ def get_upper_guards_slices(self, *, myg, mxg=0): """ xinner = self.xinner_ind if self.connection_inner_x is not None: - xinner -= mxg + xinner = xinner - mxg xouter = self.xouter_ind if self.connection_outer_x is not None: - xouter += mxg + xouter = xouter + mxg # wrap around to the beginning of the grid if necessary ystart = self.yupper_ind @@ -1189,15 +1189,15 @@ def _create_regions_toroidal(ds): # Adjust for boundary cells # keep_xboundaries is 1 if there are x-boundaries and 0 if there are not if not ds.metadata["keep_xboundaries"]: - ixs1 -= mxg - ixs2 -= mxg - nx -= 2 * mxg - jys11 += ybndry - jys21 += ybndry - ny_inner += ybndry + ybndry_upper - jys12 += ybndry + 2 * ybndry_upper - jys22 += ybndry + 2 * ybndry_upper - ny += 2 * ybndry + 2 * ybndry_upper + ixs1 = ixs1 - mxg + ixs2 = ixs2 - mxg + nx = nx - 2 * mxg + jys11 = jys11 + ybndry + jys21 = jys21 + ybndry + ny_inner = ny_inner + ybndry + ybndry_upper + jys12 = jys12 + ybndry + 2 * ybndry_upper + jys22 = jys22 + ybndry + 2 * ybndry_upper + ny = ny + 2 * ybndry + 2 * ybndry_upper # Note, include guard cells in the created regions, fill them later if topology in topologies: From 7a937179f52c484658cfe31e73a6d694d21e94dd Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 6 Jan 2023 15:38:44 +0000 Subject: [PATCH 11/34] Fix string comparison using `is` instead of `==` Can cause a bug on some systems. --- xbout/plotting/animate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index 705c386b..96a7289a 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -600,7 +600,7 @@ def animate_line( # Check plot is the right orientation t_read, x_read = data.dims - if t_read is animate_over: + if t_read == animate_over: x = x_read else: data = data.transpose(animate_over, t_read, transpose_coords=True) From 5771fe7dae68e09c15cadbaf21c4dbb88a2939b8 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 6 Jan 2023 15:29:32 +0000 Subject: [PATCH 12/34] Remove uses of deprecated Replace with a test for isinstance(..., Sequence). --- xbout/plotting/plotfuncs.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 580c9399..4a05ffb5 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence import matplotlib import matplotlib.pyplot as plt import numpy as np @@ -180,18 +181,18 @@ def plot2d_wrapper( # set up 'levels' if needed if method is xr.plot.contourf or method is xr.plot.contour: levels = kwargs.get("levels", 7) - if isinstance(levels, np.int): - levels = np.linspace(vmin, vmax, levels, endpoint=True) - # put levels back into kwargs - kwargs["levels"] = levels - else: + if isinstance(levels, Sequence): levels = np.array(list(levels)) kwargs["levels"] = levels vmin = np.min(levels) vmax = np.max(levels) + else: + levels = np.linspace(vmin, vmax, levels, endpoint=True) + # put levels back into kwargs + kwargs["levels"] = levels levels = kwargs.get("levels", 7) - if isinstance(levels, np.int): + if isinstance(levels, np.int64): levels = np.linspace(vmin, vmax, levels, endpoint=True) # put levels back into kwargs kwargs["levels"] = levels From 34cb0d06363e4a8fb4a2a32f67296c29c33c259d Mon Sep 17 00:00:00 2001 From: David Bold Date: Sat, 7 Jan 2023 19:17:15 +0100 Subject: [PATCH 13/34] Update repo data in CI --- .github/workflows/master.yml | 4 ++-- .github/workflows/pythonpackage.yml | 4 ++-- .github/workflows/pythonpublish.yml | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml index 159d0fb0..eda60534 100644 --- a/.github/workflows/master.yml +++ b/.github/workflows/master.yml @@ -33,7 +33,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - sudo apt-get install libhdf5-dev libnetcdf-dev + sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev python -m pip install --upgrade pip pip install --upgrade ${{ matrix.pip-packages }} pip install -e . @@ -66,7 +66,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - sudo apt-get install libhdf5-dev libnetcdf-dev + sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev python -m pip install --upgrade pip pip install --upgrade ${{ matrix.pip-packages }} pip install -e . diff --git a/.github/workflows/pythonpackage.yml b/.github/workflows/pythonpackage.yml index 4428378e..d07b4cc1 100644 --- a/.github/workflows/pythonpackage.yml +++ b/.github/workflows/pythonpackage.yml @@ -33,7 +33,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - sudo apt-get install libhdf5-dev libnetcdf-dev + sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev python -m pip install --upgrade pip pip install --upgrade ${{ matrix.pip-packages }} pip install -e . @@ -66,7 +66,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - sudo apt-get install libhdf5-dev libnetcdf-dev + sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev python -m pip install --upgrade pip pip install --upgrade ${{ matrix.pip-packages }} pip install -e . diff --git a/.github/workflows/pythonpublish.yml b/.github/workflows/pythonpublish.yml index 2979cc61..9f43e2bf 100644 --- a/.github/workflows/pythonpublish.yml +++ b/.github/workflows/pythonpublish.yml @@ -27,7 +27,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - sudo apt-get install libhdf5-dev libnetcdf-dev + sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev python -m pip install --upgrade pip pip install --upgrade ${{ matrix.pip-packages }} pip install -e . @@ -60,7 +60,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | - sudo apt-get install libhdf5-dev libnetcdf-dev + sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev python -m pip install --upgrade pip pip install --upgrade ${{ matrix.pip-packages }} pip install -e . @@ -123,7 +123,7 @@ jobs: python-version: '3.x' - name: Install dependencies run: | - sudo apt-get install libhdf5-dev libnetcdf-dev + sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev python -m pip install --upgrade pip pip install --upgrade setuptools wheel twine pip install -e . From 9fc746a24fb71855c6ee18ea8406d27f177029de Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 27 Mar 2020 17:33:31 +0000 Subject: [PATCH 14/34] Methods to add Cartesian (X,Y,Z) coordinates to DataArray or Dataset --- xbout/boutdataarray.py | 3 +++ xbout/boutdataset.py | 3 +++ xbout/utils.py | 21 +++++++++++++++++++++ 3 files changed, 27 insertions(+) diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index 3fc3af10..f3306098 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -385,6 +385,9 @@ def interpolate_parallel( return da + def add_cartesian_coordinates(self): + return _add_cartesian_coordinates(self.data) + def add_cartesian_coordinates(self): """ Add Cartesian (X,Y,Z) coordinates. diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 7af152fc..94ed2537 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -266,6 +266,9 @@ def find_with_dims(first_var, dims): return ds + def add_cartesian_coordinates(self): + return _add_cartesian_coordinates(self.data) + def integrate_midpoints(self, variable, *, dims=None, cumulative_t=False): """ Integrate using the midpoint rule for spatial dimensions, and trapezium rule for diff --git a/xbout/utils.py b/xbout/utils.py index 9ea1cbda..9587fcb4 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -128,6 +128,27 @@ def update_ny(name): return da +def _add_cartesian_coordinates(ds): + # Add Cartesian X and Y coordinates if they do not exist already + # Works on either BoutDataset or BoutDataArray + + R = ds["R"] + Z = ds["Z"] + zeta = ds[ds.metadata["bout_zdim"]] + if "X_cartesian" not in ds.coords: + X = R * np.cos(zeta) + ds = ds.assign_coords(X_cartesian=X) + if "Y_cartesian" not in ds.coords: + Y = R * np.sin(zeta) + ds = ds.assign_coords(Y_cartesian=Y) + if "Z_cartesian" not in ds.coords: + zcoord = ds.metadata["bout_zdim"] + nz = len(ds[zcoord]) + ds = ds.assign_coords(Z_cartesian=Z.expand_dims({zcoord: nz}, axis=-1)) + + return ds + + def _1d_coord_from_spacing(spacing, dim, ds=None, *, origin_at=None): """ Create a 1d coordinate varying along the dimension 'dim' from the grid spacing From 802e5f3dfab127b3c24d5fa7f05ecf9f1c2dd9a9 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 27 Mar 2020 17:36:08 +0000 Subject: [PATCH 15/34] plot3d() method using mayavi Port of Fabio Riva's 3d plot routine to xBOUT. --- xbout/boutdataarray.py | 3 + xbout/plotting/plotfuncs.py | 116 +++++++++++++++++++++++++++++++++++- xbout/plotting/utils.py | 62 +++++++++++++++++++ 3 files changed, 178 insertions(+), 3 deletions(-) diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index f3306098..1eace5a4 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -1080,3 +1080,6 @@ def plot_regions(self, ax=None, **kwargs): tokamak topology. """ return plotfuncs.plot_regions(self.data, ax=ax, **kwargs) + + def plot3d(self, ax=None, **kwargs): + return plotfuncs.plot3d(self.data, **kwargs) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 4a05ffb5..8b2e01ac 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -10,6 +10,8 @@ _create_norm, _decompose_regions, _is_core_only, + _k3d_plot_isel, + _make_structured_triangulation, plot_separatrices, plot_targets, ) @@ -69,7 +71,7 @@ def plot_regions(da, ax=None, **kwargs): infer_intervals=False, add_colorbar=False, ax=ax, - **kwargs + **kwargs, ) for region in colored_regions ] @@ -93,7 +95,7 @@ def plot2d_wrapper( vmax=None, aspect=None, extend=None, - **kwargs + **kwargs, ): """ Make a 2D plot using an xarray method, taking into account branch cuts (X-points). @@ -269,7 +271,7 @@ def plot2d_wrapper( add_colorbar=False, add_labels=add_label, cmap=cmap, - **kwargs + **kwargs, ) for region, add_label in zip(da_regions.values(), add_labels) ] @@ -346,3 +348,111 @@ def plot2d_wrapper( plot_targets(da_regions, ax, x=x, y=y, hatching=add_limiter_hatching) return artists + + +def plot3d(da, style="surface", engine="k3d", **kwargs): + """ + Make a 3d plot + + Parameters + ---------- + style : {'surface', 'poloidal planes'} + Type of plot to make: + - 'surface' plots the outer surface of the DataArray + - 'poloidal planes' plots each poloidal plane in the DataArray + engine : {'k3d', 'mayavi'} + 3d plotting library to use + """ + + if len(da.dims) != 3: + raise ValueError(f"plot3d needs to be passed 3d data. Got {da.dims}.") + + da = da.bout.add_cartesian_coordinates() + vmin = kwargs.pop("vmin", float(da.min().values)) + vmax = kwargs.pop("vmax", float(da.max().values)) + xcoord = da.metadata["bout_xdim"] + ycoord = da.metadata["bout_ydim"] + zcoord = da.metadata["bout_zdim"] + + if engine == "k3d": + import k3d + + plot = k3d.plot() + + for region_name, da_region in _decompose_regions(da).items(): + + npsi, ntheta, nzeta = da_region.shape + + if style == "surface": + region = da_region.regions[region_name] + + if region.connection_inner_x is None: + # Plot the inner-x surface + plot += _k3d_plot_isel(da_region, {xcoord: 0}, vmin, vmax, kwargs) + + if region.connection_outer_x is None: + # Plot the outer-x surface + plot += _k3d_plot_isel(da_region, {xcoord: -1}, vmin, vmax, kwargs) + + if region.connection_lower_y is None: + # Plot the lower-y surface + plot += _k3d_plot_isel(da_region, {ycoord: 0}, vmin, vmax, kwargs) + + if region.connection_upper_y is None: + # Plot the upper-y surface + plot += _k3d_plot_isel(da_region, {ycoord: -1}, vmin, vmax, kwargs) + + # First z-surface + plot += _k3d_plot_isel(da_region, {zcoord: 0}, vmin, vmax, kwargs) + + # Last z-surface + plot += _k3d_plot_isel(da_region, {zcoord: -1}, vmin, vmax, kwargs) + elif style == "poloidal planes": + for zeta in range(nzeta): + plot += _k3d_plot_isel( + da_region, {zcoord: zeta}, vmin, vmax, kwargs + ) + else: + raise ValueError(f"style='{style}' not implemented for engine='k3d'") + + return plot + + elif engine == "mayavi": + from mayavi import mlab + + if style == "surface": + for region_name, da_region in _decompose_regions(da).items(): + region = da_region.regions[region_name] + + # Always include z-surfaces + surface_selections = [ + {da.metadata["bout_zdim"]: 0}, + {da.metadata["bout_zdim"]: -1}, + ] + if region.connection_inner_x is None: + # Plot the inner-x surface + surface_selections.append({da.metadata["bout_xdim"]: 0}) + if region.connection_outer_x is None: + # Plot the outer-x surface + surface_selections.append({da.metadata["bout_xdim"]: -1}) + if region.connection_lower_y is None: + # Plot the lower-y surface + surface_selections.append({da.metadata["bout_ydim"]: 0}) + if region.connection_upper_y is None: + # Plot the upper-y surface + surface_selections.append({da.metadata["bout_ydim"]: -1}) + + for surface_sel in surface_selections: + da_sel = da_region.isel(surface_sel) + X = da_sel["X_cartesian"].values + Y = da_sel["Y_cartesian"].values + Z = da_sel["Z_cartesian"].values + data = da_sel.values + + mlab.mesh(X, Y, Z, scalars=data, vmin=vmin, vmax=vmax, **kwargs) + else: + raise ValueError(f"style='{style}' not implemented for engine='mayavi'") + + plt.show() + else: + raise ValueError(f"Unrecognised plot3d() 'engine' argument: {engine}") diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 0c09226e..277cd997 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -204,3 +204,65 @@ def _get_perp_vec(u1, u2, magnitude=0.04): v = np.linalg.norm((vx, vy)) wx, wy = -vy / v * magnitude, vx / v * magnitude return wx, wy + + +def _make_structured_triangulation(m, n): + # Construct a 'triangulation' of an (m x n) grid + # Returns an array of triples of (flattened) indices of the corners of the triangles + indices = np.zeros((m - 1, 2 * (n - 1), 3), dtype=np.uint32) + # indices[0] = [0, 1, n] + # indices[1] = [1, n+1, n] + + indices[:, 0 : 2 * (n - 1) : 2, 0] = ( + n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + ) + indices[:, 0 : 2 * (n - 1) : 2, 1] = ( + n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + + 1 + ) + indices[:, 0 : 2 * (n - 1) : 2, 2] = ( + n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + + n + ) + + indices[:, 1 : 2 * (n - 1) : 2, 0] = ( + n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + + 1 + ) + indices[:, 1 : 2 * (n - 1) : 2, 1] = ( + n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + + n + + 1 + ) + indices[:, 1 : 2 * (n - 1) : 2, 2] = ( + n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + + n + ) + + return indices.reshape((2 * (m - 1) * (n - 1), 3)) + + +def _k3d_plot_isel(da_region, isel, vmin, vmax, kwargs): + import k3d + + da_sel = da_region.isel(isel) + X = da_sel["X_cartesian"].values.flatten().astype(np.float32) + Y = da_sel["Y_cartesian"].values.flatten().astype(np.float32) + Z = da_sel["Z_cartesian"].values.flatten().astype(np.float32) + data = da_sel.values.flatten().astype(np.float32) + + indices = _make_structured_triangulation(*da_sel.shape) + + return k3d.mesh( + np.vstack([X, Y, Z]).T, + indices, + attribute=data, + color_range=[vmin, vmax], + **kwargs + ) From b7352b067014b3e7f4ee1107de2d8948a9f2efb3 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 2 Apr 2020 00:04:27 +0100 Subject: [PATCH 16/34] WIP: Add 'isosurface' and 'volume' 3d plot styles Currently make a plot in cartesian (R, Z, zeta) coordinates, not cylindrical. --- xbout/plotting/plotfuncs.py | 226 ++++++++++++++++++++++++++++++++++-- xbout/plotting/utils.py | 2 +- 2 files changed, 217 insertions(+), 11 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 8b2e01ac..c550a3a7 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -350,7 +350,16 @@ def plot2d_wrapper( return artists -def plot3d(da, style="surface", engine="k3d", **kwargs): +def plot3d( + da, + style="surface", + engine="k3d", + levels=None, + outputgrid=(100, 100, 25), + color_map=None, + plot=None, + **kwargs, +): """ Make a 3d plot @@ -362,6 +371,15 @@ def plot3d(da, style="surface", engine="k3d", **kwargs): - 'poloidal planes' plots each poloidal plane in the DataArray engine : {'k3d', 'mayavi'} 3d plotting library to use + levels : sequence of (float, float) + For isosurface, the pairs of (level-value, opacity) to plot + outputgrid : (int, int, int), optional + For isosurface or volume plots, the number of points to use in the Cartesian + (X,Y,Z) grid, that data is interpolated onto for plotting. + color_map : k3d color map, optional + Color map for k3d plots + plot : k3d plot instance, optional + Existing plot to add new plots to """ if len(da.dims) != 3: @@ -377,7 +395,155 @@ def plot3d(da, style="surface", engine="k3d", **kwargs): if engine == "k3d": import k3d - plot = k3d.plot() + if color_map is None: + color_map = k3d.matplotlib_color_maps.Viridis + + if plot is None: + plot = k3d.plot() + return_plot = True + else: + return_plot = False + + if style == "isosurface" or style == "volume": + data = da.copy(deep=True).load() + # data = da.bout.from_region('upper_outer_PFR').copy(deep=True).load() + datamin = data.min().item() + datamax = data.max().item() + # data[0, :, :] = datamin - 2.*(datamax - datamin) + # data[-1, :, :] = datamin - 2.*(datamax - datamin) + # data[:, 0, :] = datamin - 2.*(datamax - datamin) + # data[:, -1, :] = datamin - 2.*(datamax - datamin) + + ##interpolate to Cartesian array + # rpoints = 200 + # zpoints = 200 + # Rmin = data['R'].min() + # Rmax = data['R'].max() + # Zmin = data['Z'].min() + # Zmax = data['Z'].max() + # nx, ny, nz = data.shape + + # newR = (xr.DataArray(np.linspace(Rmin, Rmax, rpoints), dims='r') + # .expand_dims({'z': zpoints, 'zeta': nz}, axis=[0, 1])) + # newZ = (xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims='z') + # .expand_dims({'r': rpoints, 'zeta': nz}, axis=[2, 1])) + # newzeta = data['zeta'].expand_dims({'r': rpoints, 'z': zpoints}, axis=[2, 0]) + + # R = data['R'].expand_dims({'zeta': nz}, axis=2) + # Z = data['Z'].expand_dims({'zeta': nz}, axis=2) + # zeta = data['zeta'].expand_dims({'x': nx, 'theta': ny}, axis=[0, 1]) + + # from scipy.interpolate import griddata + # grid = griddata( + # (R.values.flatten(), Z.values.flatten(), + # zeta.values.flatten()), + # data.values.flatten(), + # (newR.values, newZ.values, newzeta.values), + # method='nearest') + + # if style == 'isosurface': + # plot += k3d.marching_cubes(grid.astype(np.float32), + # bounds=[Rmin, Rmax, 0., 2.*np.pi, Zmin, Zmax], + # level=1., + # ) + # elif style == 'volume': + # plot += k3d.volume(grid.astype(np.float32), + # color_range=[datamin, datamax], + # bounds=[Rmin, Rmax, 0., 2.*np.pi, Zmin, Zmax], + # ) + # else: + # raise ValueError(f'{style} not supported here') + + # return plot + + xpoints, ypoints, zpoints = outputgrid + nx, ny, nz = data.shape + + # interpolate to Cartesian array + Xmin = data["X_cartesian"].min() + Xmax = data["X_cartesian"].max() + Ymin = data["Y_cartesian"].min() + Ymax = data["Y_cartesian"].max() + Zmin = data["Z_cartesian"].min() + Zmax = data["Z_cartesian"].max() + newX = xr.DataArray(np.linspace(Xmin, Xmax, xpoints), dims="x").expand_dims( + {"y": ypoints, "z": zpoints}, axis=[1, 0] + ) + newY = xr.DataArray(np.linspace(Ymin, Ymax, ypoints), dims="y").expand_dims( + {"x": xpoints, "z": zpoints}, axis=[2, 0] + ) + newZ = xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims="z").expand_dims( + {"x": xpoints, "y": ypoints}, axis=[2, 1] + ) + newR = np.sqrt(newX**2 + newY**2) + newzeta = np.arctan2(newY, newX).values + + from scipy.interpolate import RegularGridInterpolator, griddata + + # need to create 3d arrays of x and y values + # create interpolators for x and y from R and Z + print("interpolate x") + newx = griddata( + (data["R"].values.flatten(), data["Z"].values.flatten()), + data["x"].expand_dims({"theta": ny}, axis=1).values.flatten(), + (newR.values, newZ.values), + method="linear", + ) + print("interpolate y") + newy = griddata( + (data["R"].values.flatten(), data["Z"].values.flatten()), + data["theta"].expand_dims({"x": nx}, axis=0).values.flatten(), + (newR.values, newZ.values), + method="linear", + ) + # from scipy.interpolate import griddata + # print('start interpolating') + # grid = griddata( + # (data['X_cartesian'].values.flatten(), + # data['Y_cartesian'].values.flatten(), + # data['Z_cartesian'].values.flatten()), + # data.values.flatten(), + # (newX.values, newY.values, newZ.values), + # method='nearest', + # fill_value = datamin - 2.*(datamax - datamin) + # ) + # print('done interpolating') + print("start interpolating") + x = data["x"] + y = data["theta"] + z = data["zeta"] + interp = RegularGridInterpolator( + (x, y, z), + data.values, + bounds_error=False, + fill_value=datamin - 2.0 * (datamax - datamin), + ) + grid = interp( + (newx, newy, newzeta), + method="linear", + ) + print("done interpolating") + if style == "isosurface": + if levels is None: + levels = [(0.5 * (datamin + datamax), 1.0)] + for level, opacity in levels: + plot += k3d.marching_cubes( + grid.astype(np.float32), + bounds=[Xmin, Xmax, Ymin, Ymax, Zmin, Zmax], + level=level, + color_map=color_map, + ) + elif style == "volume": + plot += k3d.volume( + grid.astype(np.float32), + color_range=[datamin, datamax], + bounds=[Xmin, Xmax, Ymin, Ymax, Zmin, Zmax], + color_map=color_map, + ) + if return_plot: + return plot + else: + return for region_name, da_region in _decompose_regions(da).items(): @@ -388,34 +554,74 @@ def plot3d(da, style="surface", engine="k3d", **kwargs): if region.connection_inner_x is None: # Plot the inner-x surface - plot += _k3d_plot_isel(da_region, {xcoord: 0}, vmin, vmax, kwargs) + plot += _k3d_plot_isel( + da_region, + {xcoord: 0}, + vmin, + vmax, + color_map=color_map, + **kwargs, + ) if region.connection_outer_x is None: # Plot the outer-x surface - plot += _k3d_plot_isel(da_region, {xcoord: -1}, vmin, vmax, kwargs) + plot += _k3d_plot_isel( + da_region, + {xcoord: -1}, + vmin, + vmax, + color_map=color_map, + **kwargs, + ) if region.connection_lower_y is None: # Plot the lower-y surface - plot += _k3d_plot_isel(da_region, {ycoord: 0}, vmin, vmax, kwargs) + plot += _k3d_plot_isel( + da_region, + {ycoord: 0}, + vmin, + vmax, + color_map=color_map, + **kwargs, + ) if region.connection_upper_y is None: # Plot the upper-y surface - plot += _k3d_plot_isel(da_region, {ycoord: -1}, vmin, vmax, kwargs) + plot += _k3d_plot_isel( + da_region, + {ycoord: -1}, + vmin, + vmax, + color_map=color_map, + **kwargs, + ) # First z-surface - plot += _k3d_plot_isel(da_region, {zcoord: 0}, vmin, vmax, kwargs) + plot += _k3d_plot_isel( + da_region, {zcoord: 0}, vmin, vmax, color_map=color_map, **kwargs + ) # Last z-surface - plot += _k3d_plot_isel(da_region, {zcoord: -1}, vmin, vmax, kwargs) + plot += _k3d_plot_isel( + da_region, {zcoord: -1}, vmin, vmax, color_map=color_map, **kwargs + ) elif style == "poloidal planes": for zeta in range(nzeta): plot += _k3d_plot_isel( - da_region, {zcoord: zeta}, vmin, vmax, kwargs + da_region, + {zcoord: zeta}, + vmin, + vmax, + color_map=color_map, + **kwargs, ) else: raise ValueError(f"style='{style}' not implemented for engine='k3d'") - return plot + if return_plot: + return plot + else: + return elif engine == "mayavi": from mayavi import mlab diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 277cd997..5a3ab7bf 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -248,7 +248,7 @@ def _make_structured_triangulation(m, n): return indices.reshape((2 * (m - 1) * (n - 1), 3)) -def _k3d_plot_isel(da_region, isel, vmin, vmax, kwargs): +def _k3d_plot_isel(da_region, isel, vmin, vmax, **kwargs): import k3d da_sel = da_region.isel(isel) From a879a892177c6253e8e3b3313040901561335f04 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 2 Apr 2020 22:11:06 +0100 Subject: [PATCH 17/34] 2-stage interpolation for 3d volume and isosurface plots Do unstructured interpolation only in 2d to get data on an R-Z grid. Do 3d interpolation with RegularGridInterpolator, which is much faster. --- xbout/plotting/plotfuncs.py | 114 +++++++++++++++++++++++++++++------- 1 file changed, 92 insertions(+), 22 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index c550a3a7..ef553170 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -466,6 +466,10 @@ def plot3d( Ymax = data["Y_cartesian"].max() Zmin = data["Z_cartesian"].min() Zmax = data["Z_cartesian"].max() + Rmin = data["R"].min() + Rmax = data["R"].max() + Zmin = data["Z"].min() + Zmax = data["Z"].max() newX = xr.DataArray(np.linspace(Xmin, Xmax, xpoints), dims="x").expand_dims( {"y": ypoints, "z": zpoints}, axis=[1, 0] ) @@ -476,26 +480,28 @@ def plot3d( {"x": xpoints, "y": ypoints}, axis=[2, 1] ) newR = np.sqrt(newX**2 + newY**2) - newzeta = np.arctan2(newY, newX).values + newzeta = np.arctan2(newY, newX) # .values - from scipy.interpolate import RegularGridInterpolator, griddata + from scipy.interpolate import ( + RegularGridInterpolator, + griddata, + LinearNDInterpolator, + ) # need to create 3d arrays of x and y values # create interpolators for x and y from R and Z - print("interpolate x") - newx = griddata( - (data["R"].values.flatten(), data["Z"].values.flatten()), - data["x"].expand_dims({"theta": ny}, axis=1).values.flatten(), - (newR.values, newZ.values), - method="linear", - ) - print("interpolate y") - newy = griddata( - (data["R"].values.flatten(), data["Z"].values.flatten()), - data["theta"].expand_dims({"x": nx}, axis=0).values.flatten(), - (newR.values, newZ.values), - method="linear", - ) + # print('interpolate x') + # newx = griddata((data['R'].values.flatten(), data['Z'].values.flatten()), + # data['x'].expand_dims({'theta': ny}, axis=1).values.flatten(), + # (newR.values, newZ.values), + # method = 'linear', + # ) + # print('interpolate y') + # newy = griddata((data['R'].values.flatten(), data['Z'].values.flatten()), + # data['theta'].expand_dims({'x': nx}, axis=0).values.flatten(), + # (newR.values, newZ.values), + # method = 'linear', + # ) # from scipy.interpolate import griddata # print('start interpolating') # grid = griddata( @@ -508,18 +514,82 @@ def plot3d( # fill_value = datamin - 2.*(datamax - datamin) # ) # print('done interpolating') + # print('start interpolating') + # x = data['x'] + # y = data['theta'] + # z = data['zeta'] + # interp = RegularGridInterpolator( + # (x, y, z), + # data.values, + # bounds_error = False, + # fill_value = datamin - 2.*(datamax - datamin) + # ) + # grid = interp((newx, newy, newzeta), + # method='linear', + # ) + # print('done interpolating') print("start interpolating") - x = data["x"] - y = data["theta"] - z = data["zeta"] + # R3d = data['R'].expand_dims({'zeta': nz}, axis=2) + # Z3d = data['Z'].expand_dims({'zeta': nz}, axis=2) + # zeta3d = data['zeta'].expand_dims({'x': nx, 'theta': ny}, axis=(0, 1)) + # grid = griddata( + # (R3d.values.flatten(), + # Z3d.values.flatten(), + # zeta3d.values.flatten()), + # data.values.flatten(), + # (newR.values, newZ.values, newzeta.values), + # method='linear', + # fill_value = datamin - 2.*(datamax - datamin) + # ) + # interp = LinearNDInterpolator((R3d.values.flatten(), Z3d.values.flatten, + # zeta3d.values.flatten()), + # data.values.flatten(), + # fill_value = datamin - 2.*(datamax - datamin) + # ) + # print('made interpolator') + # grid = interp((newR.values, newZ.values, newzeta.values), method='linear') + # Rcyl = (xr.DataArray(np.linspace(Rmin, Rmax, zpoints), dims='r') + # .expand_dims({'z': zpoints, 'zeta': nz}, axis=[1, 2])) + # Zcyl = (xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims='z') + # .expand_dims({'r': zpoints, 'zeta': nz}, axis=[0, 2])) + Rcyl = xr.DataArray( + np.linspace(Rmin, Rmax, 2 * zpoints), dims="r" + ).expand_dims({"z": 2 * zpoints}, axis=1) + Zcyl = xr.DataArray( + np.linspace(Zmin, Zmax, 2 * zpoints), dims="z" + ).expand_dims({"r": 2 * zpoints}, axis=0) + + # Interpolate in two stages for efficiency. Unstructured 3d interpolation is + # very slow. Unstructured 2d interpolation onto Cartesian (R, Z) grids, + # followed by structured 3d interpolation onto the (X, Y, Z) grid, is much + # faster. + # Structured 3d interpolation straight from (psi, theta, zeta) to (X, Y, Z) + # leaves artifacts in the output, because theta does not vary continuously + # everywhere (has branch cuts). + + # order of dimensions does not really matter here - output only depends on + # shape of newR, newZ, newzeta. Possibly more efficient to assign the 2d + # results in the loop to the last two dimensions, so put zeta first. + data_cyl = np.zeros((nz, 2 * zpoints, 2 * zpoints)) + print("interpolate poloidal planes") + for z in range(nz): + data_cyl[z] = griddata( + (data["R"].values.flatten(), data["Z"].values.flatten()), + data.isel(zeta=z).values.flatten(), + (Rcyl.values, Zcyl.values), + method="cubic", + fill_value=datamin - 2.0 * (datamax - datamin), + ) + print("build 3d interpolator") interp = RegularGridInterpolator( - (x, y, z), - data.values, + (data["zeta"].values, Rcyl.isel(z=0).values, Zcyl.isel(r=0).values), + data_cyl, bounds_error=False, fill_value=datamin - 2.0 * (datamax - datamin), ) + print("do 3d interpolation") grid = interp( - (newx, newy, newzeta), + (newzeta, newR, newZ), method="linear", ) print("done interpolating") From 3c2b12dd87fd08a92a324c013a449eb1177bf5a9 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 16 Sep 2020 22:26:36 +0100 Subject: [PATCH 18/34] Option to skip interpolation for slab simulations --- xbout/plotting/plotfuncs.py | 367 +++++++++++++++++++----------------- 1 file changed, 190 insertions(+), 177 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index ef553170..3c8a78b4 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -373,9 +373,11 @@ def plot3d( 3d plotting library to use levels : sequence of (float, float) For isosurface, the pairs of (level-value, opacity) to plot - outputgrid : (int, int, int), optional + outputgrid : (int, int, int) or None, optional For isosurface or volume plots, the number of points to use in the Cartesian - (X,Y,Z) grid, that data is interpolated onto for plotting. + (X,Y,Z) grid, that data is interpolated onto for plotting. If None, then do not + interpolate and treat "bout_xdim", "bout_ydim" and "bout_zdim" coordinates as + Cartesian (useful for slab simulations). color_map : k3d color map, optional Color map for k3d plots plot : k3d plot instance, optional @@ -414,185 +416,196 @@ def plot3d( # data[:, 0, :] = datamin - 2.*(datamax - datamin) # data[:, -1, :] = datamin - 2.*(datamax - datamin) - ##interpolate to Cartesian array - # rpoints = 200 - # zpoints = 200 - # Rmin = data['R'].min() - # Rmax = data['R'].max() - # Zmin = data['Z'].min() - # Zmax = data['Z'].max() - # nx, ny, nz = data.shape - - # newR = (xr.DataArray(np.linspace(Rmin, Rmax, rpoints), dims='r') - # .expand_dims({'z': zpoints, 'zeta': nz}, axis=[0, 1])) - # newZ = (xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims='z') - # .expand_dims({'r': rpoints, 'zeta': nz}, axis=[2, 1])) - # newzeta = data['zeta'].expand_dims({'r': rpoints, 'z': zpoints}, axis=[2, 0]) - - # R = data['R'].expand_dims({'zeta': nz}, axis=2) - # Z = data['Z'].expand_dims({'zeta': nz}, axis=2) - # zeta = data['zeta'].expand_dims({'x': nx, 'theta': ny}, axis=[0, 1]) - - # from scipy.interpolate import griddata - # grid = griddata( - # (R.values.flatten(), Z.values.flatten(), - # zeta.values.flatten()), - # data.values.flatten(), - # (newR.values, newZ.values, newzeta.values), - # method='nearest') - - # if style == 'isosurface': - # plot += k3d.marching_cubes(grid.astype(np.float32), - # bounds=[Rmin, Rmax, 0., 2.*np.pi, Zmin, Zmax], - # level=1., - # ) - # elif style == 'volume': - # plot += k3d.volume(grid.astype(np.float32), - # color_range=[datamin, datamax], - # bounds=[Rmin, Rmax, 0., 2.*np.pi, Zmin, Zmax], - # ) - # else: - # raise ValueError(f'{style} not supported here') - - # return plot - - xpoints, ypoints, zpoints = outputgrid - nx, ny, nz = data.shape - - # interpolate to Cartesian array - Xmin = data["X_cartesian"].min() - Xmax = data["X_cartesian"].max() - Ymin = data["Y_cartesian"].min() - Ymax = data["Y_cartesian"].max() - Zmin = data["Z_cartesian"].min() - Zmax = data["Z_cartesian"].max() - Rmin = data["R"].min() - Rmax = data["R"].max() - Zmin = data["Z"].min() - Zmax = data["Z"].max() - newX = xr.DataArray(np.linspace(Xmin, Xmax, xpoints), dims="x").expand_dims( - {"y": ypoints, "z": zpoints}, axis=[1, 0] - ) - newY = xr.DataArray(np.linspace(Ymin, Ymax, ypoints), dims="y").expand_dims( - {"x": xpoints, "z": zpoints}, axis=[2, 0] - ) - newZ = xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims="z").expand_dims( - {"x": xpoints, "y": ypoints}, axis=[2, 1] - ) - newR = np.sqrt(newX**2 + newY**2) - newzeta = np.arctan2(newY, newX) # .values + if outputgrid is None: + Xmin = da[da.metadata["bout_xdim"]][0] + Xmax = da[da.metadata["bout_xdim"]][-1] + Ymin = da[da.metadata["bout_ydim"]][0] + Ymax = da[da.metadata["bout_ydim"]][-1] + Zmin = da[da.metadata["bout_zdim"]][0] + Zmax = da[da.metadata["bout_zdim"]][-1] - from scipy.interpolate import ( - RegularGridInterpolator, - griddata, - LinearNDInterpolator, - ) + grid = da.astype(np.float32).values + else: + ##interpolate to Cartesian array + # rpoints = 200 + # zpoints = 200 + # Rmin = data['R'].min() + # Rmax = data['R'].max() + # Zmin = data['Z'].min() + # Zmax = data['Z'].max() + # nx, ny, nz = data.shape + + # newR = (xr.DataArray(np.linspace(Rmin, Rmax, rpoints), dims='r') + # .expand_dims({'z': zpoints, 'zeta': nz}, axis=[0, 1])) + # newZ = (xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims='z') + # .expand_dims({'r': rpoints, 'zeta': nz}, axis=[2, 1])) + # newzeta = data['zeta'].expand_dims({'r': rpoints, 'z': zpoints}, axis=[2, 0]) + + # R = data['R'].expand_dims({'zeta': nz}, axis=2) + # Z = data['Z'].expand_dims({'zeta': nz}, axis=2) + # zeta = data['zeta'].expand_dims({'x': nx, 'theta': ny}, axis=[0, 1]) + + # from scipy.interpolate import griddata + # grid = griddata( + # (R.values.flatten(), Z.values.flatten(), + # zeta.values.flatten()), + # data.values.flatten(), + # (newR.values, newZ.values, newzeta.values), + # method='nearest') + + # if style == 'isosurface': + # plot += k3d.marching_cubes(grid.astype(np.float32), + # bounds=[Rmin, Rmax, 0., 2.*np.pi, Zmin, Zmax], + # level=1., + # ) + # elif style == 'volume': + # plot += k3d.volume(grid.astype(np.float32), + # color_range=[datamin, datamax], + # bounds=[Rmin, Rmax, 0., 2.*np.pi, Zmin, Zmax], + # ) + # else: + # raise ValueError(f'{style} not supported here') + + # return plot + + xpoints, ypoints, zpoints = outputgrid + nx, ny, nz = data.shape + + # interpolate to Cartesian array + Xmin = data["X_cartesian"].min() + Xmax = data["X_cartesian"].max() + Ymin = data["Y_cartesian"].min() + Ymax = data["Y_cartesian"].max() + Zmin = data["Z_cartesian"].min() + Zmax = data["Z_cartesian"].max() + Rmin = data["R"].min() + Rmax = data["R"].max() + Zmin = data["Z"].min() + Zmax = data["Z"].max() + newX = xr.DataArray( + np.linspace(Xmin, Xmax, xpoints), dims="x" + ).expand_dims({"y": ypoints, "z": zpoints}, axis=[1, 0]) + newY = xr.DataArray( + np.linspace(Ymin, Ymax, ypoints), dims="y" + ).expand_dims({"x": xpoints, "z": zpoints}, axis=[2, 0]) + newZ = xr.DataArray( + np.linspace(Zmin, Zmax, zpoints), dims="z" + ).expand_dims({"x": xpoints, "y": ypoints}, axis=[2, 1]) + newR = np.sqrt(newX**2 + newY**2) + newzeta = np.arctan2(newY, newX) # .values + + from scipy.interpolate import ( + RegularGridInterpolator, + griddata, + LinearNDInterpolator, + ) - # need to create 3d arrays of x and y values - # create interpolators for x and y from R and Z - # print('interpolate x') - # newx = griddata((data['R'].values.flatten(), data['Z'].values.flatten()), - # data['x'].expand_dims({'theta': ny}, axis=1).values.flatten(), - # (newR.values, newZ.values), - # method = 'linear', - # ) - # print('interpolate y') - # newy = griddata((data['R'].values.flatten(), data['Z'].values.flatten()), - # data['theta'].expand_dims({'x': nx}, axis=0).values.flatten(), - # (newR.values, newZ.values), - # method = 'linear', - # ) - # from scipy.interpolate import griddata - # print('start interpolating') - # grid = griddata( - # (data['X_cartesian'].values.flatten(), - # data['Y_cartesian'].values.flatten(), - # data['Z_cartesian'].values.flatten()), - # data.values.flatten(), - # (newX.values, newY.values, newZ.values), - # method='nearest', - # fill_value = datamin - 2.*(datamax - datamin) - # ) - # print('done interpolating') - # print('start interpolating') - # x = data['x'] - # y = data['theta'] - # z = data['zeta'] - # interp = RegularGridInterpolator( - # (x, y, z), - # data.values, - # bounds_error = False, - # fill_value = datamin - 2.*(datamax - datamin) - # ) - # grid = interp((newx, newy, newzeta), - # method='linear', - # ) - # print('done interpolating') - print("start interpolating") - # R3d = data['R'].expand_dims({'zeta': nz}, axis=2) - # Z3d = data['Z'].expand_dims({'zeta': nz}, axis=2) - # zeta3d = data['zeta'].expand_dims({'x': nx, 'theta': ny}, axis=(0, 1)) - # grid = griddata( - # (R3d.values.flatten(), - # Z3d.values.flatten(), - # zeta3d.values.flatten()), - # data.values.flatten(), - # (newR.values, newZ.values, newzeta.values), - # method='linear', - # fill_value = datamin - 2.*(datamax - datamin) - # ) - # interp = LinearNDInterpolator((R3d.values.flatten(), Z3d.values.flatten, - # zeta3d.values.flatten()), - # data.values.flatten(), - # fill_value = datamin - 2.*(datamax - datamin) - # ) - # print('made interpolator') - # grid = interp((newR.values, newZ.values, newzeta.values), method='linear') - # Rcyl = (xr.DataArray(np.linspace(Rmin, Rmax, zpoints), dims='r') - # .expand_dims({'z': zpoints, 'zeta': nz}, axis=[1, 2])) - # Zcyl = (xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims='z') - # .expand_dims({'r': zpoints, 'zeta': nz}, axis=[0, 2])) - Rcyl = xr.DataArray( - np.linspace(Rmin, Rmax, 2 * zpoints), dims="r" - ).expand_dims({"z": 2 * zpoints}, axis=1) - Zcyl = xr.DataArray( - np.linspace(Zmin, Zmax, 2 * zpoints), dims="z" - ).expand_dims({"r": 2 * zpoints}, axis=0) - - # Interpolate in two stages for efficiency. Unstructured 3d interpolation is - # very slow. Unstructured 2d interpolation onto Cartesian (R, Z) grids, - # followed by structured 3d interpolation onto the (X, Y, Z) grid, is much - # faster. - # Structured 3d interpolation straight from (psi, theta, zeta) to (X, Y, Z) - # leaves artifacts in the output, because theta does not vary continuously - # everywhere (has branch cuts). - - # order of dimensions does not really matter here - output only depends on - # shape of newR, newZ, newzeta. Possibly more efficient to assign the 2d - # results in the loop to the last two dimensions, so put zeta first. - data_cyl = np.zeros((nz, 2 * zpoints, 2 * zpoints)) - print("interpolate poloidal planes") - for z in range(nz): - data_cyl[z] = griddata( - (data["R"].values.flatten(), data["Z"].values.flatten()), - data.isel(zeta=z).values.flatten(), - (Rcyl.values, Zcyl.values), - method="cubic", + # need to create 3d arrays of x and y values + # create interpolators for x and y from R and Z + # print('interpolate x') + # newx = griddata((data['R'].values.flatten(), data['Z'].values.flatten()), + # data['x'].expand_dims({'theta': ny}, axis=1).values.flatten(), + # (newR.values, newZ.values), + # method = 'linear', + # ) + # print('interpolate y') + # newy = griddata((data['R'].values.flatten(), data['Z'].values.flatten()), + # data['theta'].expand_dims({'x': nx}, axis=0).values.flatten(), + # (newR.values, newZ.values), + # method = 'linear', + # ) + # from scipy.interpolate import griddata + # print('start interpolating') + # grid = griddata( + # (data['X_cartesian'].values.flatten(), + # data['Y_cartesian'].values.flatten(), + # data['Z_cartesian'].values.flatten()), + # data.values.flatten(), + # (newX.values, newY.values, newZ.values), + # method='nearest', + # fill_value = datamin - 2.*(datamax - datamin) + # ) + # print('done interpolating') + # print('start interpolating') + # x = data['x'] + # y = data['theta'] + # z = data['zeta'] + # interp = RegularGridInterpolator( + # (x, y, z), + # data.values, + # bounds_error = False, + # fill_value = datamin - 2.*(datamax - datamin) + # ) + # grid = interp((newx, newy, newzeta), + # method='linear', + # ) + # print('done interpolating') + print("start interpolating") + # R3d = data['R'].expand_dims({'zeta': nz}, axis=2) + # Z3d = data['Z'].expand_dims({'zeta': nz}, axis=2) + # zeta3d = data['zeta'].expand_dims({'x': nx, 'theta': ny}, axis=(0, 1)) + # grid = griddata( + # (R3d.values.flatten(), + # Z3d.values.flatten(), + # zeta3d.values.flatten()), + # data.values.flatten(), + # (newR.values, newZ.values, newzeta.values), + # method='linear', + # fill_value = datamin - 2.*(datamax - datamin) + # ) + # interp = LinearNDInterpolator((R3d.values.flatten(), Z3d.values.flatten, + # zeta3d.values.flatten()), + # data.values.flatten(), + # fill_value = datamin - 2.*(datamax - datamin) + # ) + # print('made interpolator') + # grid = interp((newR.values, newZ.values, newzeta.values), method='linear') + # Rcyl = (xr.DataArray(np.linspace(Rmin, Rmax, zpoints), dims='r') + # .expand_dims({'z': zpoints, 'zeta': nz}, axis=[1, 2])) + # Zcyl = (xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims='z') + # .expand_dims({'r': zpoints, 'zeta': nz}, axis=[0, 2])) + Rcyl = xr.DataArray( + np.linspace(Rmin, Rmax, 2 * zpoints), dims="r" + ).expand_dims({"z": 2 * zpoints}, axis=1) + Zcyl = xr.DataArray( + np.linspace(Zmin, Zmax, 2 * zpoints), dims="z" + ).expand_dims({"r": 2 * zpoints}, axis=0) + + # Interpolate in two stages for efficiency. Unstructured 3d interpolation is + # very slow. Unstructured 2d interpolation onto Cartesian (R, Z) grids, + # followed by structured 3d interpolation onto the (X, Y, Z) grid, is much + # faster. + # Structured 3d interpolation straight from (psi, theta, zeta) to (X, Y, Z) + # leaves artifacts in the output, because theta does not vary continuously + # everywhere (has branch cuts). + + # order of dimensions does not really matter here - output only depends on + # shape of newR, newZ, newzeta. Possibly more efficient to assign the 2d + # results in the loop to the last two dimensions, so put zeta first. + data_cyl = np.zeros((nz, 2 * zpoints, 2 * zpoints)) + print("interpolate poloidal planes") + for z in range(nz): + data_cyl[z] = griddata( + (data["R"].values.flatten(), data["Z"].values.flatten()), + data.isel(zeta=z).values.flatten(), + (Rcyl.values, Zcyl.values), + method="cubic", + fill_value=datamin - 2.0 * (datamax - datamin), + ) + print("build 3d interpolator") + interp = RegularGridInterpolator( + (data["zeta"].values, Rcyl.isel(z=0).values, Zcyl.isel(r=0).values), + data_cyl, + bounds_error=False, fill_value=datamin - 2.0 * (datamax - datamin), ) - print("build 3d interpolator") - interp = RegularGridInterpolator( - (data["zeta"].values, Rcyl.isel(z=0).values, Zcyl.isel(r=0).values), - data_cyl, - bounds_error=False, - fill_value=datamin - 2.0 * (datamax - datamin), - ) - print("do 3d interpolation") - grid = interp( - (newzeta, newR, newZ), - method="linear", - ) - print("done interpolating") + print("do 3d interpolation") + grid = interp( + (newzeta, newR, newZ), + method="linear", + ) + print("done interpolating") + if style == "isosurface": if levels is None: levels = [(0.5 * (datamin + datamax), 1.0)] From 1116029aad2a124807a3c83b222f7c90c78bb228 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 9 Dec 2022 17:23:12 +0000 Subject: [PATCH 19/34] Arguments for extra mayavi settings in plot3d() --- xbout/plotting/plotfuncs.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 3c8a78b4..2e7264e0 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -358,6 +358,9 @@ def plot3d( outputgrid=(100, 100, 25), color_map=None, plot=None, + mayavi_figure=None, + mayavi_figure_args=None, + mayavi_view=None, **kwargs, ): """ @@ -382,6 +385,15 @@ def plot3d( Color map for k3d plots plot : k3d plot instance, optional Existing plot to add new plots to + mayavi_figure : mayavi.core.scene.Scene, default None + Existing Mayavi figure to add this plot to. + mayavi_figure_args : dict, default None + Arguments to use when creating a new Mayavi figure. Ignored if `mayavi_figure` + is passed. + mayavi_view : (float, float, float), default None + If set, arguments are passed to mlab.view() to set the view when engine="mayavi" + **kwargs + Extra keyword arguments are passed to the backend plotting function """ if len(da.dims) != 3: @@ -709,6 +721,16 @@ def plot3d( elif engine == "mayavi": from mayavi import mlab + if mayavi_figure is None: + if mayavi_figure_args is None: + mayavi_figure_args = {} + mlab.figure(**mayavi_figure_args) + else: + mlab.figure(mayavi_figure) + + if mayavi_view is not None: + mlab.view(*mayavi_view) + if style == "surface": for region_name, da_region in _decompose_regions(da).items(): region = da_region.regions[region_name] From 4a9666e839b52770ebf01b54c23e91f85463c784 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 9 Dec 2022 17:25:41 +0000 Subject: [PATCH 20/34] Allow passing indices to select different surfaces in plot3d() --- xbout/plotting/plotfuncs.py | 33 +++++++++++++++++++++++++++------ 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 2e7264e0..f01a67ec 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -358,6 +358,9 @@ def plot3d( outputgrid=(100, 100, 25), color_map=None, plot=None, + surface_xinds=None, + surface_yinds=None, + surface_zinds=None, mayavi_figure=None, mayavi_figure_args=None, mayavi_view=None, @@ -385,6 +388,18 @@ def plot3d( Color map for k3d plots plot : k3d plot instance, optional Existing plot to add new plots to + surface_xinds : (int, int), default None + Indices to select when plotting radial surfaces. These indices are local to the + region being plotted, so values will be strange. Recommend using values relative + to the radial boundaries (i.e. positive for inner boundary and negative for + outer boundary). + surface_yinds : (int, int), default None + Indices to select when plotting poloidal surfaces. These indices are local to the + region being plotted, so values will be strange. Recommend using values relative + to the poloidal boundaries (i.e. positive for lower boundaries and negative for + upper boundaries). + surface_zinds : (int, int), default None + Indices to select when plotting toroidal surfaces mayavi_figure : mayavi.core.scene.Scene, default None Existing Mayavi figure to add this plot to. mayavi_figure_args : dict, default None @@ -736,22 +751,28 @@ def plot3d( region = da_region.regions[region_name] # Always include z-surfaces + zstart_ind = 0 if surface_zinds is None else surface_zinds[0] + zend_ind = -1 if surface_zinds is None else surface_zinds[1] surface_selections = [ - {da.metadata["bout_zdim"]: 0}, - {da.metadata["bout_zdim"]: -1}, + {da.metadata["bout_zdim"]: zstart_ind}, + {da.metadata["bout_zdim"]: zend_ind}, ] if region.connection_inner_x is None: # Plot the inner-x surface - surface_selections.append({da.metadata["bout_xdim"]: 0}) + xstart_ind = 0 if surface_xinds is None else surface_xinds[0] + surface_selections.append({xcoord: xstart_ind}) if region.connection_outer_x is None: # Plot the outer-x surface - surface_selections.append({da.metadata["bout_xdim"]: -1}) + xend_ind = -1 if surface_xinds is None else surface_xinds[1] + surface_selections.append({xcoord: xend_ind}) if region.connection_lower_y is None: # Plot the lower-y surface - surface_selections.append({da.metadata["bout_ydim"]: 0}) + ystart_ind = 0 if surface_yinds is None else surface_yinds[0] + surface_selections.append({ycoord: ystart_ind}) if region.connection_upper_y is None: # Plot the upper-y surface - surface_selections.append({da.metadata["bout_ydim"]: -1}) + yend_ind = -1 if surface_yinds is None else surface_yinds[1] + surface_selections.append({ycoord: yend_ind}) for surface_sel in surface_selections: da_sel = da_region.isel(surface_sel) From 7e61dd6a7585b26297b8fc3bc7d715ba0025c038 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 9 Dec 2022 18:25:39 +0000 Subject: [PATCH 21/34] Option to animate mayavi figures If DataArray with time dimension is passed, an animation will be made. --- xbout/plotting/plotfuncs.py | 137 +++++++++++++++++++++++++----------- 1 file changed, 97 insertions(+), 40 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index f01a67ec..c635bf9a 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -411,17 +411,40 @@ def plot3d( Extra keyword arguments are passed to the backend plotting function """ - if len(da.dims) != 3: - raise ValueError(f"plot3d needs to be passed 3d data. Got {da.dims}.") - - da = da.bout.add_cartesian_coordinates() - vmin = kwargs.pop("vmin", float(da.min().values)) - vmax = kwargs.pop("vmax", float(da.max().values)) + tcoord = da.metadata["bout_tdim"] xcoord = da.metadata["bout_xdim"] ycoord = da.metadata["bout_ydim"] zcoord = da.metadata["bout_zdim"] + if tcoord in da.dims: + animate = True + if len(da.dims) != 4: + raise ValueError( + f"plot3d needs to be passed 3d spatial data. Got {da.dims}." + ) + else: + animate = False + if len(da.dims) != 3: + raise ValueError( + f"plot3d needs to be passed 3d spatial data. Got {da.dims}." + ) + + da = da.bout.add_cartesian_coordinates() + if "vmin" in kwargs: + vmin = kwargs.pop("vmin") + else: + vmin = float(da.min().values) + if "vmax" in kwargs: + vmax = kwargs.pop("vmax") + else: + vmax = float(da.max().values) + if engine == "k3d": + if animate: + raise ValueError( + "animation not supported by k3d, do not pass time-dependent DataArray" + ) + import k3d if color_map is None: @@ -747,41 +770,75 @@ def plot3d( mlab.view(*mayavi_view) if style == "surface": - for region_name, da_region in _decompose_regions(da).items(): - region = da_region.regions[region_name] - - # Always include z-surfaces - zstart_ind = 0 if surface_zinds is None else surface_zinds[0] - zend_ind = -1 if surface_zinds is None else surface_zinds[1] - surface_selections = [ - {da.metadata["bout_zdim"]: zstart_ind}, - {da.metadata["bout_zdim"]: zend_ind}, - ] - if region.connection_inner_x is None: - # Plot the inner-x surface - xstart_ind = 0 if surface_xinds is None else surface_xinds[0] - surface_selections.append({xcoord: xstart_ind}) - if region.connection_outer_x is None: - # Plot the outer-x surface - xend_ind = -1 if surface_xinds is None else surface_xinds[1] - surface_selections.append({xcoord: xend_ind}) - if region.connection_lower_y is None: - # Plot the lower-y surface - ystart_ind = 0 if surface_yinds is None else surface_yinds[0] - surface_selections.append({ycoord: ystart_ind}) - if region.connection_upper_y is None: - # Plot the upper-y surface - yend_ind = -1 if surface_yinds is None else surface_yinds[1] - surface_selections.append({ycoord: yend_ind}) - - for surface_sel in surface_selections: - da_sel = da_region.isel(surface_sel) - X = da_sel["X_cartesian"].values - Y = da_sel["Y_cartesian"].values - Z = da_sel["Z_cartesian"].values - data = da_sel.values - mlab.mesh(X, Y, Z, scalars=data, vmin=vmin, vmax=vmax, **kwargs) + def create_or_update_plot(plot_objects=None, tind=None): + if plot_objects is None: + # Creating plot for first time + plot_objects_to_return = {} + this_da = da + if tind is not None: + this_da = da.isel({tcoord: tind}) + + for region_name, da_region in _decompose_regions(this_da).items(): + region = da_region.regions[region_name] + + # Always include z-surfaces + zstart_ind = 0 if surface_zinds is None else surface_zinds[0] + zend_ind = -1 if surface_zinds is None else surface_zinds[1] + surface_selections = [ + {this_da.metadata["bout_zdim"]: zstart_ind}, + {this_da.metadata["bout_zdim"]: zend_ind}, + ] + if region.connection_inner_x is None: + # Plot the inner-x surface + xstart_ind = 0 if surface_xinds is None else surface_xinds[0] + surface_selections.append({xcoord: xstart_ind}) + if region.connection_outer_x is None: + # Plot the outer-x surface + xend_ind = -1 if surface_xinds is None else surface_xinds[1] + surface_selections.append({xcoord: xend_ind}) + if region.connection_lower_y is None: + # Plot the lower-y surface + ystart_ind = 0 if surface_yinds is None else surface_yinds[0] + surface_selections.append({ycoord: ystart_ind}) + if region.connection_upper_y is None: + # Plot the upper-y surface + yend_ind = -1 if surface_yinds is None else surface_yinds[1] + surface_selections.append({ycoord: yend_ind}) + + for i, surface_sel in enumerate(surface_selections): + da_sel = da_region.isel(surface_sel) + X = da_sel["X_cartesian"].values + Y = da_sel["Y_cartesian"].values + Z = da_sel["Z_cartesian"].values + data = da_sel.values + + if plot_objects is None: + plot_objects_to_return[region_name + str(i)] = mlab.mesh( + X, Y, Z, scalars=data, vmin=vmin, vmax=vmax, **kwargs + ) + else: + plot_objects[ + region_name + str(i) + ].mlab_source.scalars = data + if plot_objects is None: + return plot_objects_to_return + + if animate: + print(f"tind=0") + plot_objects = create_or_update_plot(tind=0) + + @mlab.animate + def animation_func(): + for tind in range(1, da.sizes[tcoord]): + print(f"tind={tind}") + create_or_update_plot(plot_objects=plot_objects, tind=tind) + yield + + a = animation_func() + mlab.show() + else: + create_or_update_plot() else: raise ValueError(f"style='{style}' not implemented for engine='mayavi'") From a1d8c49364a6af7d744dcb0f5d75638c8710f18e Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 9 Dec 2022 18:26:04 +0000 Subject: [PATCH 22/34] Option to save mayavi figures If animation is being made, will be saved as a sequence of files. --- xbout/plotting/plotfuncs.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index c635bf9a..696c8875 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -358,6 +358,7 @@ def plot3d( outputgrid=(100, 100, 25), color_map=None, plot=None, + save_as=None, surface_xinds=None, surface_yinds=None, surface_zinds=None, @@ -388,6 +389,9 @@ def plot3d( Color map for k3d plots plot : k3d plot instance, optional Existing plot to add new plots to + save_as : str + Filename to save figure to. Animations will be saved as a sequence of + numbered files. surface_xinds : (int, int), default None Indices to select when plotting radial surfaces. These indices are local to the region being plotted, so values will be strange. Recommend using values relative @@ -444,6 +448,8 @@ def plot3d( raise ValueError( "animation not supported by k3d, do not pass time-dependent DataArray" ) + if save_as is not None: + raise ValueError("save_as not supported by k3d implementation yet") import k3d @@ -821,6 +827,14 @@ def create_or_update_plot(plot_objects=None, tind=None): plot_objects[ region_name + str(i) ].mlab_source.scalars = data + if save_as: + if tind is None: + mlab.savefig(save_as) + else: + name_parts = save_as.split(".") + name_parts = name_parts[:-1] + [str(tind)] + name_parts[-1:] + this_save_as = ".".join(name_parts) + mlab.savefig(this_save_as) if plot_objects is None: return plot_objects_to_return From 3905400297cf520fbd81fa8e51cf14ccca9a9352 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 10 Dec 2022 11:40:40 +0000 Subject: [PATCH 23/34] Fix use of mlab.view(*mayavi_view) mlab.view() must be called after the plot is created in order for the `focalpoint` argument to work. --- xbout/plotting/plotfuncs.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 696c8875..d8f07057 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -772,9 +772,6 @@ def plot3d( else: mlab.figure(mayavi_figure) - if mayavi_view is not None: - mlab.view(*mayavi_view) - if style == "surface": def create_or_update_plot(plot_objects=None, tind=None): @@ -827,6 +824,9 @@ def create_or_update_plot(plot_objects=None, tind=None): plot_objects[ region_name + str(i) ].mlab_source.scalars = data + if mayavi_view is not None: + mlab.view(*mayavi_view) + if save_as: if tind is None: mlab.savefig(save_as) @@ -853,6 +853,7 @@ def animation_func(): mlab.show() else: create_or_update_plot() + else: raise ValueError(f"style='{style}' not implemented for engine='mayavi'") From 65d8cd985ac7f1fab45e86c8e4be1eaded34c874 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 10 Dec 2022 11:42:03 +0000 Subject: [PATCH 24/34] Option for colorbar in 3d plots So far only implemented for mayavi. --- xbout/plotting/plotfuncs.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index d8f07057..550ba437 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -357,6 +357,8 @@ def plot3d( levels=None, outputgrid=(100, 100, 25), color_map=None, + colorbar=True, + colorbar_font_size=None, plot=None, save_as=None, surface_xinds=None, @@ -387,6 +389,11 @@ def plot3d( Cartesian (useful for slab simulations). color_map : k3d color map, optional Color map for k3d plots + colorbar : bool or dict, default True + Add a color bar. If a dict is passed, it is passed on to the colorbar() function + as keyword arguments. + colorbar_font_size : float, default None + Set the font size used by the colorbar (for engine="mayavi") plot : k3d plot instance, optional Existing plot to add new plots to save_as : str @@ -450,6 +457,8 @@ def plot3d( ) if save_as is not None: raise ValueError("save_as not supported by k3d implementation yet") + if colorbar: + warnings.warn("colorbar not added to k3d plots yet") import k3d @@ -824,9 +833,21 @@ def create_or_update_plot(plot_objects=None, tind=None): plot_objects[ region_name + str(i) ].mlab_source.scalars = data + if mayavi_view is not None: mlab.view(*mayavi_view) + if colorbar and (tind is None or tind == 0): + if isinstance(colorbar, dict): + colorbar_args = colorbar + else: + colorbar_args = {} + cb = mlab.colorbar(**colorbar_args) + if colorbar_font_size is not None: + cb.scalar_bar.unconstrained_font_size = True + cb.label_text_property.font_size = colorbar_font_size + cb.title_text_property.font_size = colorbar_font_size + if save_as: if tind is None: mlab.savefig(save_as) From 65ebfe1a357f9b4b028faff2d452d309dde18e0f Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 6 Jan 2023 14:21:37 +0000 Subject: [PATCH 25/34] Mayavi animation creates .gif directly ...rather than a collection of .png files. Uses `wand`, which is a Python interface to ImageMagick. Also tweaks the vmin/vmax handling in plot3d() so that vmin=None and vmax=None can be passed explicitly to have vmin/vmax be evaluated for each frame individually in an animation. --- xbout/plotting/plotfuncs.py | 94 +++++++++++++++++++++++++++++-------- 1 file changed, 75 insertions(+), 19 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 550ba437..520f8d7e 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -2,6 +2,8 @@ import matplotlib import matplotlib.pyplot as plt import numpy as np +from pathlib import Path +from tempfile import TemporaryDirectory import warnings import xarray as xr @@ -364,6 +366,7 @@ def plot3d( surface_xinds=None, surface_yinds=None, surface_zinds=None, + fps=20.0, mayavi_figure=None, mayavi_figure_args=None, mayavi_view=None, @@ -411,6 +414,8 @@ def plot3d( upper boundaries). surface_zinds : (int, int), default None Indices to select when plotting toroidal surfaces + fps : float, default 20 + Frames per second to use when creating an animation. mayavi_figure : mayavi.core.scene.Scene, default None Existing Mayavi figure to add this plot to. mayavi_figure_args : dict, default None @@ -418,6 +423,11 @@ def plot3d( is passed. mayavi_view : (float, float, float), default None If set, arguments are passed to mlab.view() to set the view when engine="mayavi" + vmin, vmax : float + vmin and vmax are treated specially. If a float is passed, then it is used for + vmin/vmax. If the arguments are not passed, then the minimum and maximum of the + data are used. For an animation, to get minimum and/or maximum calculated + separately for each frame, pass `vmin=None` and/or `vmax=None` explicitly. **kwargs Extra keyword arguments are passed to the backend plotting function """ @@ -783,7 +793,7 @@ def plot3d( if style == "surface": - def create_or_update_plot(plot_objects=None, tind=None): + def create_or_update_plot(plot_objects=None, tind=None, this_save_as=None): if plot_objects is None: # Creating plot for first time plot_objects_to_return = {} @@ -848,32 +858,78 @@ def create_or_update_plot(plot_objects=None, tind=None): cb.label_text_property.font_size = colorbar_font_size cb.title_text_property.font_size = colorbar_font_size - if save_as: + if this_save_as: if tind is None: - mlab.savefig(save_as) + mlab.savefig(this_save_as) else: - name_parts = save_as.split(".") + name_parts = this_save_as.split(".") name_parts = name_parts[:-1] + [str(tind)] + name_parts[-1:] - this_save_as = ".".join(name_parts) - mlab.savefig(this_save_as) + frame_save_as = ".".join(name_parts) + mlab.savefig(frame_save_as) if plot_objects is None: return plot_objects_to_return if animate: - print(f"tind=0") - plot_objects = create_or_update_plot(tind=0) - - @mlab.animate - def animation_func(): - for tind in range(1, da.sizes[tcoord]): - print(f"tind={tind}") - create_or_update_plot(plot_objects=plot_objects, tind=tind) - yield - - a = animation_func() - mlab.show() + orig_offscreen_option = mlab.options.offscreen + mlab.options.offscreen = True + + try: + # resets mlab.options.offscreen when it finishes, even if there is + # an error + if save_as is None: + raise ValueError( + "Must pass `save_as` for a mayavi animation, or no output will " + "be created" + ) + with TemporaryDirectory() as d: + nframes = da.sizes[tcoord] + + # First create png files in the temporary directory + temp_path = Path(d) + temp_save_as = str(temp_path.joinpath("temp.png")) + print(f"tind=0") + plot_objects = create_or_update_plot( + tind=0, this_save_as=temp_save_as + ) + + # @mlab.animate # interative mayavi animation too slow + def animation_func(): + for tind in range(1, nframes): + print(f"tind={tind}") + create_or_update_plot( + plot_objects=plot_objects, + tind=tind, + this_save_as=temp_save_as, + ) + # yield # needed for an interactive mayavi animation + + a = animation_func() + + # Use ImageMagick via the wand package to turn the .png files into + # an animation + try: + from wand.image import Image + except ImportError: + raise ImportError( + "Please install the `wand` package to save the 3d animation" + ) + with Image() as animation: + for i in range(nframes): + filename = str(temp_path.joinpath(f"temp.{i}.png")) + with Image(filename=filename) as frame: + animation.sequence.append(frame) + animation.type = "optimize" + animation.loop = 0 + # Delay is in milliseconds + animation.delay = int(round(1000.0 / fps)) + animation.save(filename=save_as) + finally: + mlab.options.offscreen = orig_offscreen_option + + # mlab.show() # interative mayavi animation so slow that it's not useful + return a else: - create_or_update_plot() + create_or_update_plot(this_save_as=save_as) else: raise ValueError(f"style='{style}' not implemented for engine='mayavi'") From 2348dc92d2796da71df4279a7d54cd8656845470 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 6 Jan 2023 14:33:56 +0000 Subject: [PATCH 26/34] Add warning to plot3d() docstrings --- xbout/boutdataarray.py | 14 ++++++++++++++ xbout/plotting/plotfuncs.py | 6 ++++++ 2 files changed, 20 insertions(+) diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index 1eace5a4..7e57cf1e 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -1082,4 +1082,18 @@ def plot_regions(self, ax=None, **kwargs): return plotfuncs.plot_regions(self.data, ax=ax, **kwargs) def plot3d(self, ax=None, **kwargs): + """ + Make a 3d plot + + Warnings + -------- + + 3d plotting functionality is still a bit of a work in progress. Bugs are likely, and + help developing is welcome! + + Parameters + ---------- + + See plotfuncs.plot3d() + """ return plotfuncs.plot3d(self.data, **kwargs) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 520f8d7e..31498ccf 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -375,6 +375,12 @@ def plot3d( """ Make a 3d plot + Warnings + -------- + + 3d plotting functionality is still a bit of a work in progress. Bugs are likely, and + help developing is welcome! + Parameters ---------- style : {'surface', 'poloidal planes'} From a5eb843340789c7ceb04b26957a6cf6b94fa1275 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 7 Jan 2023 18:56:55 +0000 Subject: [PATCH 27/34] Simplify slicing and reshape in _make_structured_triangulation() Co-authored-by: David Bold --- xbout/plotting/utils.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 5a3ab7bf..5f2ab14f 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -213,39 +213,39 @@ def _make_structured_triangulation(m, n): # indices[0] = [0, 1, n] # indices[1] = [1, n+1, n] - indices[:, 0 : 2 * (n - 1) : 2, 0] = ( + indices[:, ::2, 0] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] ) - indices[:, 0 : 2 * (n - 1) : 2, 1] = ( + indices[:, ::2, 1] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + 1 ) - indices[:, 0 : 2 * (n - 1) : 2, 2] = ( + indices[:, ::2, 2] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + n ) - indices[:, 1 : 2 * (n - 1) : 2, 0] = ( + indices[:, 1::2, 0] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + 1 ) - indices[:, 1 : 2 * (n - 1) : 2, 1] = ( + indices[:, 1::2, 1] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + n + 1 ) - indices[:, 1 : 2 * (n - 1) : 2, 2] = ( + indices[:, 1::2, 2] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + n ) - return indices.reshape((2 * (m - 1) * (n - 1), 3)) + return indices.reshape((-1, 3)) def _k3d_plot_isel(da_region, isel, vmin, vmax, **kwargs): From 953b8c49d30017eaaaa7ebdcde1dab9b1dea1253 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sat, 7 Jan 2023 19:07:40 +0000 Subject: [PATCH 28/34] Better comments in _make_structured_triangulation() --- xbout/plotting/utils.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 5f2ab14f..218c903a 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -213,32 +213,44 @@ def _make_structured_triangulation(m, n): # indices[0] = [0, 1, n] # indices[1] = [1, n+1, n] + # Each quadrilateral grid cell with corners (i,j), (i,j+1), (i+1,j+1), (i+1,j), at + # flattened index I=i*n+j is split into two triangles, one with corners (i,j), + # (i,j+1), (i+1,j) and the other with corners (i,j+1) (i+1,j+1), (i+1,j), + + # Lower left triangles + # (i,j) corner indices[:, ::2, 0] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] ) + # (i,j+1) corner indices[:, ::2, 1] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + 1 ) + # (i+1,j) corner indices[:, ::2, 2] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + n ) + # Upper right triangles + # (i,j+1) corner indices[:, 1::2, 0] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + 1 ) + # (i+1,j+1) corner indices[:, 1::2, 1] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] + n + 1 ) + # (i+1,j) corner indices[:, 1::2, 2] = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] From 2b2ded5e6ce4c49b92777e09d5cd7ff497da2d4f Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 8 Jan 2023 11:59:21 +0000 Subject: [PATCH 29/34] Better import handling for optional k3d and mayavi packages --- xbout/plotting/plotfuncs.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 31498ccf..0569fb86 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -476,7 +476,12 @@ def plot3d( if colorbar: warnings.warn("colorbar not added to k3d plots yet") - import k3d + try: + import k3d + except ImportError: + raise ImportError( + 'Please install the `k3d` package for 3d plotting with `engine="k3d"`' + ) if color_map is None: color_map = k3d.matplotlib_color_maps.Viridis @@ -788,7 +793,13 @@ def plot3d( return elif engine == "mayavi": - from mayavi import mlab + try: + from mayavi import mlab + except ImportError: + raise ImportError( + "Please install the `mayavi` package for 3d plotting with " + '`engine="mayavi"`' + ) if mayavi_figure is None: if mayavi_figure_args is None: From 0f189a2f919826701d2d8497c6cbfa01817642b7 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 8 Jan 2023 12:02:09 +0000 Subject: [PATCH 30/34] Tidy up index creation in _make_structured_triangulation() --- xbout/plotting/utils.py | 38 +++++++++----------------------------- 1 file changed, 9 insertions(+), 29 deletions(-) diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index 218c903a..e94e2f8c 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -217,45 +217,25 @@ def _make_structured_triangulation(m, n): # flattened index I=i*n+j is split into two triangles, one with corners (i,j), # (i,j+1), (i+1,j) and the other with corners (i,j+1) (i+1,j+1), (i+1,j), - # Lower left triangles - # (i,j) corner - indices[:, ::2, 0] = ( + lower_left_indices = ( n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] ) + # Lower left triangles + # (i,j) corner + indices[:, ::2, 0] = lower_left_indices # (i,j+1) corner - indices[:, ::2, 1] = ( - n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] - + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] - + 1 - ) + indices[:, ::2, 1] = lower_left_indices + 1 # (i+1,j) corner - indices[:, ::2, 2] = ( - n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] - + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] - + n - ) + indices[:, ::2, 2] = lower_left_indices + n # Upper right triangles # (i,j+1) corner - indices[:, 1::2, 0] = ( - n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] - + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] - + 1 - ) + indices[:, 1::2, 0] = lower_left_indices + 1 # (i+1,j+1) corner - indices[:, 1::2, 1] = ( - n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] - + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] - + n - + 1 - ) + indices[:, 1::2, 1] = lower_left_indices + n + 1 # (i+1,j) corner - indices[:, 1::2, 2] = ( - n * np.arange(m - 1, dtype=np.uint32)[:, np.newaxis] - + np.arange((n - 1), dtype=np.uint32)[np.newaxis, :] - + n - ) + indices[:, 1::2, 2] = lower_left_indices + n return indices.reshape((-1, 3)) From 4c42a4ff9d8a4e5f09d706acd18a700579974ba1 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 8 Jan 2023 12:05:00 +0000 Subject: [PATCH 31/34] Remove commented out interpolation code in plot3d() Commented out code was from earlier, less successful, attempts at interpolation. --- xbout/plotting/plotfuncs.py | 110 ------------------------------------ 1 file changed, 110 deletions(-) diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 0569fb86..732d818e 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -494,13 +494,8 @@ def plot3d( if style == "isosurface" or style == "volume": data = da.copy(deep=True).load() - # data = da.bout.from_region('upper_outer_PFR').copy(deep=True).load() datamin = data.min().item() datamax = data.max().item() - # data[0, :, :] = datamin - 2.*(datamax - datamin) - # data[-1, :, :] = datamin - 2.*(datamax - datamin) - # data[:, 0, :] = datamin - 2.*(datamax - datamin) - # data[:, -1, :] = datamin - 2.*(datamax - datamin) if outputgrid is None: Xmin = da[da.metadata["bout_xdim"]][0] @@ -512,48 +507,6 @@ def plot3d( grid = da.astype(np.float32).values else: - ##interpolate to Cartesian array - # rpoints = 200 - # zpoints = 200 - # Rmin = data['R'].min() - # Rmax = data['R'].max() - # Zmin = data['Z'].min() - # Zmax = data['Z'].max() - # nx, ny, nz = data.shape - - # newR = (xr.DataArray(np.linspace(Rmin, Rmax, rpoints), dims='r') - # .expand_dims({'z': zpoints, 'zeta': nz}, axis=[0, 1])) - # newZ = (xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims='z') - # .expand_dims({'r': rpoints, 'zeta': nz}, axis=[2, 1])) - # newzeta = data['zeta'].expand_dims({'r': rpoints, 'z': zpoints}, axis=[2, 0]) - - # R = data['R'].expand_dims({'zeta': nz}, axis=2) - # Z = data['Z'].expand_dims({'zeta': nz}, axis=2) - # zeta = data['zeta'].expand_dims({'x': nx, 'theta': ny}, axis=[0, 1]) - - # from scipy.interpolate import griddata - # grid = griddata( - # (R.values.flatten(), Z.values.flatten(), - # zeta.values.flatten()), - # data.values.flatten(), - # (newR.values, newZ.values, newzeta.values), - # method='nearest') - - # if style == 'isosurface': - # plot += k3d.marching_cubes(grid.astype(np.float32), - # bounds=[Rmin, Rmax, 0., 2.*np.pi, Zmin, Zmax], - # level=1., - # ) - # elif style == 'volume': - # plot += k3d.volume(grid.astype(np.float32), - # color_range=[datamin, datamax], - # bounds=[Rmin, Rmax, 0., 2.*np.pi, Zmin, Zmax], - # ) - # else: - # raise ValueError(f'{style} not supported here') - - # return plot - xpoints, ypoints, zpoints = outputgrid nx, ny, nz = data.shape @@ -586,70 +539,7 @@ def plot3d( LinearNDInterpolator, ) - # need to create 3d arrays of x and y values - # create interpolators for x and y from R and Z - # print('interpolate x') - # newx = griddata((data['R'].values.flatten(), data['Z'].values.flatten()), - # data['x'].expand_dims({'theta': ny}, axis=1).values.flatten(), - # (newR.values, newZ.values), - # method = 'linear', - # ) - # print('interpolate y') - # newy = griddata((data['R'].values.flatten(), data['Z'].values.flatten()), - # data['theta'].expand_dims({'x': nx}, axis=0).values.flatten(), - # (newR.values, newZ.values), - # method = 'linear', - # ) - # from scipy.interpolate import griddata - # print('start interpolating') - # grid = griddata( - # (data['X_cartesian'].values.flatten(), - # data['Y_cartesian'].values.flatten(), - # data['Z_cartesian'].values.flatten()), - # data.values.flatten(), - # (newX.values, newY.values, newZ.values), - # method='nearest', - # fill_value = datamin - 2.*(datamax - datamin) - # ) - # print('done interpolating') - # print('start interpolating') - # x = data['x'] - # y = data['theta'] - # z = data['zeta'] - # interp = RegularGridInterpolator( - # (x, y, z), - # data.values, - # bounds_error = False, - # fill_value = datamin - 2.*(datamax - datamin) - # ) - # grid = interp((newx, newy, newzeta), - # method='linear', - # ) - # print('done interpolating') print("start interpolating") - # R3d = data['R'].expand_dims({'zeta': nz}, axis=2) - # Z3d = data['Z'].expand_dims({'zeta': nz}, axis=2) - # zeta3d = data['zeta'].expand_dims({'x': nx, 'theta': ny}, axis=(0, 1)) - # grid = griddata( - # (R3d.values.flatten(), - # Z3d.values.flatten(), - # zeta3d.values.flatten()), - # data.values.flatten(), - # (newR.values, newZ.values, newzeta.values), - # method='linear', - # fill_value = datamin - 2.*(datamax - datamin) - # ) - # interp = LinearNDInterpolator((R3d.values.flatten(), Z3d.values.flatten, - # zeta3d.values.flatten()), - # data.values.flatten(), - # fill_value = datamin - 2.*(datamax - datamin) - # ) - # print('made interpolator') - # grid = interp((newR.values, newZ.values, newzeta.values), method='linear') - # Rcyl = (xr.DataArray(np.linspace(Rmin, Rmax, zpoints), dims='r') - # .expand_dims({'z': zpoints, 'zeta': nz}, axis=[1, 2])) - # Zcyl = (xr.DataArray(np.linspace(Zmin, Zmax, zpoints), dims='z') - # .expand_dims({'r': zpoints, 'zeta': nz}, axis=[0, 2])) Rcyl = xr.DataArray( np.linspace(Rmin, Rmax, 2 * zpoints), dims="r" ).expand_dims({"z": 2 * zpoints}, axis=1) From e18370b9d708dd3243df517ecf743eefcb9ccf20 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 8 Jan 2023 12:10:33 +0000 Subject: [PATCH 32/34] Add k3d and mayavi to [options.extras_require] in setup.cfg Minimum versions are a guess - haven't actually tested that they do work or that earlier versions would not be compatible. --- setup.cfg | 2 ++ 1 file changed, 2 insertions(+) diff --git a/setup.cfg b/setup.cfg index 1f6349b5..300e8534 100644 --- a/setup.cfg +++ b/setup.cfg @@ -50,6 +50,8 @@ calc = xrft xhistogram docs = sphinx >= 1.4 +3d_plot_k3d = k3d >= 2.8.0 +3d_plot_mayavi = mayavi >= 4.7.2 [build_sphinx] project = $metadata.name From b9823248b52fcec6f961a8f78a199d2825381a55 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 9 Jan 2023 10:46:05 +0000 Subject: [PATCH 33/34] Collect 3d plot dependencies in single optional dependency Cleaner than separate meta-packages and no one is forced to use it anyway, and the packages give helpful error messages if they are not installed. Co-authored-by: David Bold --- setup.cfg | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/setup.cfg b/setup.cfg index 300e8534..a2d71621 100644 --- a/setup.cfg +++ b/setup.cfg @@ -50,8 +50,10 @@ calc = xrft xhistogram docs = sphinx >= 1.4 -3d_plot_k3d = k3d >= 2.8.0 -3d_plot_mayavi = mayavi >= 4.7.2 +3d_plot = + k3d >= 2.8.0 + mayavi >= 4.7.2 + wand [build_sphinx] project = $metadata.name From 41b31cfa318852e960d1e7f4cc7fb4d934b01774 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 9 Jan 2023 12:22:35 +0000 Subject: [PATCH 34/34] Remove Python-3.7, add 3.10 in 'publish' workflow Matches regular CI workflow. [skip ci] --- .github/workflows/pythonpublish.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pythonpublish.yml b/.github/workflows/pythonpublish.yml index 9f43e2bf..44dbd492 100644 --- a/.github/workflows/pythonpublish.yml +++ b/.github/workflows/pythonpublish.yml @@ -14,7 +14,7 @@ jobs: if: always() strategy: matrix: - python-version: [3.7, 3.8, 3.9] + python-version: [3.8, 3.9, '3.10'] pip-packages: - "setuptools pip pytest pytest-cov coverage codecov boutdata xarray numpy>=1.16.0" fail-fast: true @@ -47,7 +47,7 @@ jobs: if: always() strategy: matrix: - python-version: [3.7, 3.8] + python-version: [3.8] pip-packages: - "setuptools pip pytest pytest-cov coverage codecov boutdata==0.1.4 xarray==0.18.0 dask==2.10.0 numpy==1.18.0 natsort==5.5.0 matplotlib==3.1.1 animatplot==0.4.2 netcdf4==1.4.2 Pillow==7.2.0" # test with oldest supported version of packages. Note, using numpy==1.18.0 as a workaround because numpy==1.17.0 is not supported on Python-3.7, even though we should currently support numpy==1.17.0. fail-fast: true