diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 07d5117..56e5f3e 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -1,7 +1,7 @@ # This workflow will install Python dependencies, run tests and lint with a variety of Python versions # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python -name: pism-ragis +name: pypism on: push: diff --git a/notebooks/test_extract_profile.ipynb b/notebooks/test_extract_profile.ipynb index 9818c56..09ade4e 100644 --- a/notebooks/test_extract_profile.ipynb +++ b/notebooks/test_extract_profile.ipynb @@ -2,7 +2,19 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": 15, + "id": "1944771d", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "import xarray as xr" + ] + }, + { + "cell_type": "code", + "execution_count": 16, "id": "b2105341", "metadata": {}, "outputs": [], @@ -90,23 +102,10 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 17, "id": "a641f41f", "metadata": {}, - "outputs": [ - { - "ename": "ModuleNotFoundError", - "evalue": "No module named 'netCDF4'", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[2], line 13\u001b[0m\n\u001b[1;32m 10\u001b[0m fd, filename \u001b[38;5;241m=\u001b[39m tempfile\u001b[38;5;241m.\u001b[39mmkstemp(suffix\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m.nc\u001b[39m\u001b[38;5;124m\"\u001b[39m, prefix\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mextract_profile_test_\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 11\u001b[0m os\u001b[38;5;241m.\u001b[39mclose(fd)\n\u001b[0;32m---> 13\u001b[0m \u001b[43mcreate_dummy_input_file\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilename\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mF\u001b[49m\u001b[43m)\u001b[49m\n", - "Cell \u001b[0;32mIn[1], line 5\u001b[0m, in \u001b[0;36mcreate_dummy_input_file\u001b[0;34m(filename, F)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcreate_dummy_input_file\u001b[39m(filename, F):\n\u001b[1;32m 2\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Create an input file for testing. Does not use unlimited\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;124;03m dimensions, creates one time record only.\"\"\"\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnetCDF4\u001b[39;00m\n\u001b[1;32m 7\u001b[0m nc \u001b[38;5;241m=\u001b[39m netCDF4\u001b[38;5;241m.\u001b[39mDataset(filename, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mw\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 9\u001b[0m Mx \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m88\u001b[39m\n", - "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'netCDF4'" - ] - } - ], + "outputs": [], "source": [ " def F(x, y, z):\n", " \"\"\"A function linear in x, y, and z. Used to test our interpolation\n", @@ -120,13 +119,2709 @@ " fd, filename = tempfile.mkstemp(suffix=\".nc\", prefix=\"extract_profile_test_\")\n", " os.close(fd)\n", "\n", - " create_dummy_input_file(filename, F)\n" + " create_dummy_input_file(filename, F)\n", + " create_dummy_input_file(\"test.nc\", F)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 199, + "id": "73156724", + "metadata": {}, + "outputs": [], + "source": [ + "def create_dummy_input_file_xr(filename, F):\n", + " \"\"\"Create an input file for testing. Does not use unlimited\n", + " dimensions, creates one time record only.\"\"\"\n", + "\n", + " Mx = 88\n", + " My = 152\n", + " Mz = 11\n", + " Mt = 1\n", + " \n", + " # use X and Y ranges corresponding to a grid covering Greenland\n", + " x = np.linspace(-669650.0, 896350.0, Mx)\n", + " y = np.linspace(-3362600.0, -644600.0, My)\n", + " z = np.linspace(0, 4000.0, Mz)\n", + "\n", + " xx, yy = np.meshgrid(x, y)\n", + "\n", + " def write(prefix, dimensions):\n", + " \"Write test data to the file using given storage order.\"\n", + "\n", + " slices = {\"x\": slice(0, Mx), \"y\": slice(0, My), \"time\": 0, \"z\": None}\n", + " dim_map = {\"x\": Mx, \"y\": My, \"z\": Mz, \"time\": Mt}\n", + "\n", + " # set indexes for all dimensions (z index will be re-set below)\n", + " indexes = [Ellipsis] * len(dimensions)\n", + " for k, d in enumerate(dimensions):\n", + " indexes[k] = slices[d]\n", + "\n", + " # transpose 2D array if needed\n", + " if dimensions.index(\"y\") < dimensions.index(\"x\"):\n", + "\n", + " def T(x):\n", + " return x\n", + "\n", + " else:\n", + " T = np.transpose\n", + "\n", + " dims = [dim_map[d] for d in dimensions]\n", + " variable = np.zeros(dims)\n", + " if \"z\" in dimensions:\n", + " for k in range(Mz):\n", + " indexes[dimensions.index(\"z\")] = k\n", + " variable[*indexes] = T(F(xx, yy, z[k]))\n", + " else:\n", + " variable[*indexes] = T(F(xx, yy, 0))\n", + " \n", + " return (dimensions, variable, {\"long_name\": name + \" (make it long!)\"})\n", + "\n", + " from itertools import permutations\n", + "\n", + " def P(x):\n", + " return list(permutations(x))\n", + "\n", + " data_vars = dict()\n", + " for d in sorted(P([\"x\", \"y\"]) + P([\"time\", \"x\", \"y\"])):\n", + " prefix = \"test_2D_\"\n", + " name = prefix + \"_\".join(d)\n", + " data_vars[name] = write(prefix, d)\n", + "\n", + " for d in sorted(P([\"x\", \"y\", \"z\"]) + P([\"time\", \"x\", \"y\", \"z\"])):\n", + " prefix = \"test_3D_\"\n", + " name = prefix + \"_\".join(d)\n", + " data_vars[name] = write(prefix, d)\n", + " \n", + " ds = xr.Dataset(\n", + " data_vars=data_vars,\n", + " coords=dict(\n", + " time=([\"time\"], [0], {}),\n", + " z=([\"z\"], z, {\"_FillValue\": False, \"units\": \"m\"}),\n", + " y=([\"y\"], y, {\"_FillValue\": False, \"units\": \"m\", \"axis\": \"Y\", \"standard_name\": \"projection_y_coordinate\"}),\n", + " x=([\"x\"], x, {\"_FillValue\": False, \"units\": \"m\", \"axis\": \"X\", \"standard_name\": \"projection_x_coordinate\"}),\n", + " ),\n", + " attrs=dict(description=\"Test data.\"),\n", + " )\n", + " ds.to_netcdf(filename)\n", + " return ds\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 200, + "id": "1acde809", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.Dataset>\n",
+       "Dimensions:             (time: 1, x: 88, y: 152, z: 11)\n",
+       "Coordinates:\n",
+       "  * time                (time) int64 0\n",
+       "  * z                   (z) float64 0.0 400.0 800.0 ... 3.2e+03 3.6e+03 4e+03\n",
+       "  * y                   (y) float64 -3.363e+06 -3.345e+06 ... -6.446e+05\n",
+       "  * x                   (x) float64 -6.696e+05 -6.516e+05 ... 8.964e+05\n",
+       "Data variables: (12/38)\n",
+       "    test_2D_time_x_y    (time, x, y) float64 -7.394e+04 ... -3.918e+03\n",
+       "    test_2D_time_y_x    (time, y, x) float64 -7.394e+04 ... -3.918e+03\n",
+       "    test_2D_x_time_y    (x, time, y) float64 -7.394e+04 ... -3.918e+03\n",
+       "    test_2D_x_y         (x, y) float64 -7.394e+04 -7.358e+04 ... -3.918e+03\n",
+       "    test_2D_x_y_time    (x, y, time) float64 -7.394e+04 ... -3.918e+03\n",
+       "    test_2D_y_time_x    (y, time, x) float64 -7.394e+04 ... -3.918e+03\n",
+       "    ...                  ...\n",
+       "    test_3D_z_x_time_y  (z, x, time, y) float64 -7.394e+04 ... -3.758e+03\n",
+       "    test_3D_z_x_y       (z, x, y) float64 -7.394e+04 -7.358e+04 ... -3.758e+03\n",
+       "    test_3D_z_x_y_time  (z, x, y, time) float64 -7.394e+04 ... -3.758e+03\n",
+       "    test_3D_z_y_time_x  (z, y, time, x) float64 -7.394e+04 ... -3.758e+03\n",
+       "    test_3D_z_y_x       (z, y, x) float64 -7.394e+04 -7.376e+04 ... -3.758e+03\n",
+       "    test_3D_z_y_x_time  (z, y, x, time) float64 -7.394e+04 ... -3.758e+03\n",
+       "Attributes:\n",
+       "    description:  Test data.
" + ], + "text/plain": [ + "\n", + "Dimensions: (time: 1, x: 88, y: 152, z: 11)\n", + "Coordinates:\n", + " * time (time) int64 0\n", + " * z (z) float64 0.0 400.0 800.0 ... 3.2e+03 3.6e+03 4e+03\n", + " * y (y) float64 -3.363e+06 -3.345e+06 ... -6.446e+05\n", + " * x (x) float64 -6.696e+05 -6.516e+05 ... 8.964e+05\n", + "Data variables: (12/38)\n", + " test_2D_time_x_y (time, x, y) float64 -7.394e+04 ... -3.918e+03\n", + " test_2D_time_y_x (time, y, x) float64 -7.394e+04 ... -3.918e+03\n", + " test_2D_x_time_y (x, time, y) float64 -7.394e+04 ... -3.918e+03\n", + " test_2D_x_y (x, y) float64 -7.394e+04 -7.358e+04 ... -3.918e+03\n", + " test_2D_x_y_time (x, y, time) float64 -7.394e+04 ... -3.918e+03\n", + " test_2D_y_time_x (y, time, x) float64 -7.394e+04 ... -3.918e+03\n", + " ... ...\n", + " test_3D_z_x_time_y (z, x, time, y) float64 -7.394e+04 ... -3.758e+03\n", + " test_3D_z_x_y (z, x, y) float64 -7.394e+04 -7.358e+04 ... -3.758e+03\n", + " test_3D_z_x_y_time (z, x, y, time) float64 -7.394e+04 ... -3.758e+03\n", + " test_3D_z_y_time_x (z, y, time, x) float64 -7.394e+04 ... -3.758e+03\n", + " test_3D_z_y_x (z, y, x) float64 -7.394e+04 -7.376e+04 ... -3.758e+03\n", + " test_3D_z_y_x_time (z, y, x, time) float64 -7.394e+04 ... -3.758e+03\n", + "Attributes:\n", + " description: Test data." + ] + }, + "execution_count": 200, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "create_dummy_input_file_xr(\"test_xr.nc\", F)" + ] + }, + { + "cell_type": "code", + "execution_count": 85, + "id": "c0e0227f", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "Mx = 88\n", + "My = 152\n", + "Mz = 11\n", + "\n", + "slices = {\"x\": slice(0, Mx), \"y\": slice(0, My), \"time\": 0, \"z\": None}" + ] + }, + { + "cell_type": "code", + "execution_count": 54, + "id": "24bda91c", + "metadata": {}, + "outputs": [], + "source": [ + "a =slices[\"y\"]" + ] + }, + { + "cell_type": "code", + "execution_count": 48, + "id": "a4c6f8fb", + "metadata": {}, + "outputs": [ + { + "ename": "TypeError", + "evalue": "expected a sequence of integers or a single integer, got '{'x': slice(0, 88, None), 'y': slice(0, 152, None), 'time': 0, 'z': None}'", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[48], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mzeros\u001b[49m\u001b[43m(\u001b[49m\u001b[43mslices\u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[0;31mTypeError\u001b[0m: expected a sequence of integers or a single integer, got '{'x': slice(0, 88, None), 'y': slice(0, 152, None), 'time': 0, 'z': None}'" + ] + } + ], + "source": [ + "np.zeros(slices)" + ] + }, + { + "cell_type": "code", + "execution_count": 67, + "id": "33a70c4d", + "metadata": {}, + "outputs": [], + "source": [ + "dimension = ('time', 'x', 'y')\n", + "dim_map = {\"x\": Mx, \"y\": My, \"z\": None, \"time\": 1}\n" + ] + }, + { + "cell_type": "code", + "execution_count": 57, + "id": "aee3b59b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'x': slice(0, 88, None), 'y': slice(0, 152, None), 'time': 0, 'z': None}" + ] + }, + "execution_count": 57, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "slices" + ] + }, + { + "cell_type": "code", + "execution_count": 58, + "id": "2ee394cf", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'time'" + ] + }, + "execution_count": 58, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "d[0]" + ] + }, + { + "cell_type": "code", + "execution_count": 66, + "id": "c914b489", + "metadata": {}, + "outputs": [ + { + "ename": "KeyError", + "evalue": "('time', 'x', 'y')", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mKeyError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[66], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m [\u001b[38;5;241m*\u001b[39m(\u001b[43mdims\u001b[49m\u001b[43m[\u001b[49m\u001b[43md\u001b[49m\u001b[43m]\u001b[49m)]\n", + "\u001b[0;31mKeyError\u001b[0m: ('time', 'x', 'y')" + ] + } + ], + "source": [ + "[*(dims[d])]" + ] + }, + { + "cell_type": "code", + "execution_count": 71, + "id": "da8b27bd", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(1, 88, 152)" + ] + }, + "execution_count": 71, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "np.zeros([dim_map[d] for d in dimension]).shape" + ] + }, + { + "cell_type": "code", + "execution_count": 80, + "id": "199c14fe", + "metadata": {}, + "outputs": [], + "source": [ + "variable = np.zeros([dim_map[d] for d in dimension])\n" + ] + }, + { + "cell_type": "code", + "execution_count": 82, + "id": "bc775fca", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([[0., 0., 0., ..., 0., 0., 0.],\n", + " [0., 0., 0., ..., 0., 0., 0.],\n", + " [0., 0., 0., ..., 0., 0., 0.],\n", + " ...,\n", + " [0., 0., 0., ..., 0., 0., 0.],\n", + " [0., 0., 0., ..., 0., 0., 0.],\n", + " [0., 0., 0., ..., 0., 0., 0.]])" + ] + }, + "execution_count": 82, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "variable[0, slice(0, 88, None), slice(0, 152, None)]" + ] + }, + { + "cell_type": "code", + "execution_count": 91, + "id": "226d49a5", + "metadata": {}, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'ds' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[91], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mds\u001b[49m\n", + "\u001b[0;31mNameError\u001b[0m: name 'ds' is not defined" + ] + } + ], + "source": [ + "ds" + ] + }, + { + "cell_type": "code", + "execution_count": 95, + "id": "ec3b1f65", + "metadata": {}, + "outputs": [], + "source": [ + " \"\"\"\n", + " Return a dataset with missing values\n", + " \"\"\"\n", + "\n", + " start_date = \"1980-01-01\"\n", + " calendar = \"standard\"\n", + " periodicity = \"YS\"\n", + " reference_time = pd.Timestamp(\"1980-01-01\")\n", + "\n", + " nx = 101\n", + " ny = 101\n", + " nt = 2\n", + "\n", + " dates = xr.cftime_range(start_date, periods=nt, freq=periodicity, calendar=calendar)\n", + " x = np.linspace(-100_000, 100_000, nx)\n", + " y = np.linspace(-100_000, 100_000, ny)\n", + " t = np.arange(nt)\n", + " X, Y = np.meshgrid(x, y)\n", + "\n", + " t_mask = ((np.abs(X)<25_000) & (np.abs(Y)<25_000)).reshape(1, ny, nx)\n", + " t_mask = np.repeat(t_mask, repeats=nt, axis=0)\n", + " temperature = np.exp(((X/100_000)**2 + (Y/100_000)**2)).reshape(1, ny, nx)\n", + " temperature = np.repeat(temperature, repeats=nt, axis=0)\n", + "\n", + " p_mask = ((np.abs(X)<25_000) & (Y<-25_000)).reshape(1, ny, nx)\n", + " p_mask = np.repeat(p_mask, repeats=nt, axis=0)\n", + " precipitation = (np.exp(((X/100_000)**2 + (Y/100_000)**2)) ** 2).reshape(1, ny, nx)\n", + " precipitation = np.repeat(precipitation, repeats=nt, axis=0)\n", + "\n", + " ds_true = xr.Dataset(\n", + " data_vars=dict(\n", + " temperature=([\"time\", \"x\", \"y\"], temperature, {\"units\": \"Celsius\"}),\n", + " precipitation=([\"time\", \"y\", \"x\"], precipitation, {\"units\": \"mm/day\"}),\n", + " ),\n", + " coords=dict(\n", + " time=dates,\n", + " y=([\"y\"], y, {\"_FillValue\": False, \"units\": \"m\", \"axis\": \"Y\", \"standard_name\": \"projection_y_coordinate\"}),\n", + " x=([\"x\"], x, {\"_FillValue\": False, \"units\": \"m\", \"axis\": \"X\", \"standard_name\": \"projection_x_coordinate\"}),\n", + " reference_time=reference_time,\n", + " ),\n", + " attrs=dict(description=\"Test data.\"),\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": 96, + "id": "3659a8f6", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.Dataset>\n",
+       "Dimensions:         (time: 2, x: 101, y: 101)\n",
+       "Coordinates:\n",
+       "  * time            (time) object 1980-01-01 00:00:00 1981-01-01 00:00:00\n",
+       "  * y               (y) float64 -1e+05 -9.8e+04 -9.6e+04 ... 9.8e+04 1e+05\n",
+       "  * x               (x) float64 -1e+05 -9.8e+04 -9.6e+04 ... 9.8e+04 1e+05\n",
+       "    reference_time  datetime64[ns] 1980-01-01\n",
+       "Data variables:\n",
+       "    temperature     (time, x, y) float64 7.389 7.102 6.832 ... 6.832 7.102 7.389\n",
+       "    precipitation   (time, y, x) float64 54.6 50.44 46.67 ... 46.67 50.44 54.6\n",
+       "Attributes:\n",
+       "    description:  Test data.
" + ], + "text/plain": [ + "\n", + "Dimensions: (time: 2, x: 101, y: 101)\n", + "Coordinates:\n", + " * time (time) object 1980-01-01 00:00:00 1981-01-01 00:00:00\n", + " * y (y) float64 -1e+05 -9.8e+04 -9.6e+04 ... 9.8e+04 1e+05\n", + " * x (x) float64 -1e+05 -9.8e+04 -9.6e+04 ... 9.8e+04 1e+05\n", + " reference_time datetime64[ns] 1980-01-01\n", + "Data variables:\n", + " temperature (time, x, y) float64 7.389 7.102 6.832 ... 6.832 7.102 7.389\n", + " precipitation (time, y, x) float64 54.6 50.44 46.67 ... 46.67 50.44 54.6\n", + "Attributes:\n", + " description: Test data." + ] + }, + "execution_count": 96, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ds_true" ] }, { "cell_type": "code", "execution_count": null, - "id": "1944771d", + "id": "e6e2fad4", "metadata": {}, "outputs": [], "source": [] @@ -148,7 +2843,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.6" + "version": "3.11.5" } }, "nbformat": 4, diff --git a/tests/test_extract_profile.py b/tests/test_extract_profile.py new file mode 100644 index 0000000..9de5199 --- /dev/null +++ b/tests/test_extract_profile.py @@ -0,0 +1,130 @@ +# Copyright (C) 2015, 2016, 2018, 2021, 2023 Constantine Khroulev and Andy Aschwanden +# +# This file is part of pypism. +# +# PISM-RAGIS is free software; you can redistribute it and/or modify it under the +# terms of the GNU General Public License as published by the Free Software +# Foundation; either version 3 of the License, or (at your option) any later +# version. +# +# PISM-RAGIS is distributed in the hope that it will be useful, but WITHOUT ANY +# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS +# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more +# details. +# +# You should have received a copy of the GNU General Public License +# along with PISM; if not, write to the Free Software +# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA + +""" +Test extract_profile +""" + +import numpy as np +import pytest +import xarray as xr + + +def F(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray: + """A function linear in x, y, and z. Used to test our interpolation + scheme.""" + return 10.0 + 0.01 * x + 0.02 * y + 0.03 + 0.04 * z + + +def create_dummy_input_file(F) -> xr.Dataset: + """Create an input file for testing. Does not use unlimited + dimensions, creates one time record only.""" + + Mx = 88 + My = 152 + Mz = 11 + Mt = 1 + + # use X and Y ranges corresponding to a grid covering Greenland + x = np.linspace(-669650.0, 896350.0, Mx) + y = np.linspace(-3362600.0, -644600.0, My) + z = np.linspace(0, 4000.0, Mz) + + xx, yy = np.meshgrid(x, y) + + def write(prefix, dimensions): + "Write test data to the file using given storage order." + + slices = {"x": slice(0, Mx), "y": slice(0, My), "time": 0, "z": None} + dim_map = {"x": Mx, "y": My, "z": Mz, "time": Mt} + + # set indexes for all dimensions (z index will be re-set below) + indexes = [Ellipsis] * len(dimensions) + for k, d in enumerate(dimensions): + indexes[k] = slices[d] + + # transpose 2D array if needed + if dimensions.index("y") < dimensions.index("x"): + + def T(x): + return x + + else: + T = np.transpose + + dims = [dim_map[d] for d in dimensions] + variable = np.zeros(dims) + if "z" in dimensions: + for k in range(Mz): + indexes[dimensions.index("z")] = k + variable[*indexes] = T(F(xx, yy, z[k])) + else: + variable[*indexes] = T(F(xx, yy, 0)) + + return (dimensions, variable, {"long_name": name + " (make it long!)"}) + + from itertools import permutations + + def P(x): + return list(permutations(x)) + + data_vars = dict() + for d in sorted(P(["x", "y"]) + P(["time", "x", "y"])): + prefix = "test_2D_" + name = prefix + "_".join(d) + data_vars[name] = write(prefix, d) + + for d in sorted(P(["x", "y", "z"]) + P(["time", "x", "y", "z"])): + prefix = "test_3D_" + name = prefix + "_".join(d) + data_vars[name] = write(prefix, d) + + ds = xr.Dataset( + data_vars=data_vars, + coords=dict( + time=(["time"], [0], {}), + z=(["z"], z, {"_FillValue": False, "units": "m"}), + y=( + ["y"], + y, + { + "_FillValue": False, + "units": "m", + "axis": "Y", + "standard_name": "projection_y_coordinate", + }, + ), + x=( + ["x"], + x, + { + "_FillValue": False, + "units": "m", + "axis": "X", + "standard_name": "projection_x_coordinate", + }, + ), + ), + attrs=dict(description="Test data."), + ) + return ds + + +@pytest.fixture(name="create_dummpy_input_file") +def fixture_create_dummy_input_file_xyz(): + return create_dummy_input_file(F)