Allow changing scales during field getitem/setitem

Author: kburns

dedalus/core/field.py

@@ -575,13 +575,23 @@ def __init__(self, dist, bases=None, name=None, tensorsig=None, dtype=None):
         # Add weak reference to distributor
         dist.fields.add(self)
 
-    def __getitem__(self, layout):
-        """Return data viewed in specified layout."""
+    def __getitem__(self, key):
+        """Return data viewed in specified layout and scales."""
+        if isinstance(key, tuple):
+            layout, scales = key
+            self.change_scales(scales)
+        else:
+            layout = key
         self.change_layout(layout)
         return self.data
 
-    def __setitem__(self, layout, data):
-        """Set data viewed in a specified layout."""
+    def __setitem__(self, key, data):
+        """Set data viewed in a specified layout and scales."""
+        if isinstance(key, tuple):
+            layout, scales = key
+            self.preset_scales(scales)
+        else:
+            layout = key
         layout = self.dist.get_layout_object(layout)
         self.preset_layout(layout)
         copyto(self.data, data)

docs/notebooks/dedalus_tutorial_2.ipynb

@@ -212,7 +212,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/var/folders/gl/8q1_pm2s1490lvyfvm_8yby80000gn/T/ipykernel_4766/1812536267.py:4: RuntimeWarning: divide by zero encountered in log10\n",
+      "/var/folders/gl/8q1_pm2s1490lvyfvm_8yby80000gn/T/ipykernel_26772/1812536267.py:4: RuntimeWarning: divide by zero encountered in log10\n",
       "  log_mag = lambda xmesh, ymesh, data: (xmesh, ymesh, np.log10(np.abs(data)))\n"
      ]
     },
@@ -351,6 +351,32 @@
     "plot_bot_2d(f, title=\"f['g']\", figkw=figkw);"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Scale factors can also be provided as second arguments when setting or retrieving field data via the geitem/setitem interface:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(64, 64)\n",
+      "(128, 128)\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(f['g', 1].shape)\n",
+    "print(f['g', 2].shape)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -378,7 +404,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
@@ -405,7 +431,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
@@ -448,7 +474,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
@@ -473,7 +499,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [
     {
@@ -506,7 +532,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
@@ -550,7 +576,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [
     {
@@ -581,7 +607,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
@@ -625,7 +651,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -649,7 +675,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [
     {
@@ -694,7 +720,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -718,7 +744,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
@@ -757,7 +783,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [
     {

docs/notebooks/dedalus_tutorial_3.ipynb

@@ -54,16 +54,14 @@
 solver.stop_sim_time = stop_sim_time
 
 # Main loop
-u.change_scales(1)
-u_list = [np.copy(u['g'])]
+u_list = [u['g',1].copy()]
 t_list = [solver.sim_time]
 while solver.proceed:
     solver.step(timestep)
     if solver.iteration % 100 == 0:
         logger.info('Iteration=%i, Time=%e, dt=%e' %(solver.iteration, solver.sim_time, timestep))
     if solver.iteration % 25 == 0:
-        u.change_scales(1)
-        u_list.append(np.copy(u['g']))
+        u_list.append(u['g',1].copy())
         t_list.append(solver.sim_time)
 
 # Plot

examples/ivp_1d_kdv_burgers/kdv_burgers.py

@@ -103,7 +103,6 @@
 flow.add_property(np.sqrt(u@u)/nu, name='Re')
 
 # Main loop
-startup_iter = 10
 try:
     logger.info('Starting main loop')
     while solver.proceed:
@@ -117,4 +116,3 @@
     raise
 finally:
     solver.log_stats()
-

examples/ivp_2d_rayleigh_benard/rayleigh_benard.py

@@ -79,16 +79,15 @@
 # Solver
 solver = problem.build_solver(ncc_cutoff=ncc_cutoff)
 pert_norm = np.inf
-f.change_scales(dealias)
-steps = [f['g'].ravel().copy()]
+steps = [f['g',1].copy().ravel()]
 while pert_norm > tolerance:
     solver.newton_iteration()
     pert_norm = sum(pert.allreduce_data_norm('c', 2) for pert in solver.perturbations)
     logger.info(f'Perturbation norm: {pert_norm:.3e}')
     f0 = f(r=0).evaluate().allgather_data('g')[0,0,0]
     Ri = f0**((n-1)/2)
     logger.info(f'R iterate: {Ri}')
-    steps.append(f['g'].ravel().copy())
+    steps.append(f['g',1].copy().ravel())
 
 # Compare to reference solutions from Boyd
 R_ref = {0.0: np.sqrt(6),
@@ -110,7 +109,6 @@
 
 # Plot solution
 plt.figure(figsize=(6, 4))
-_, _, r = dist.local_grids(ball, scales=(dealias,dealias,dealias))
 alpha = np.linspace(0.2, 1, len(steps))
 color = ('C0',) * (len(steps)-1) + ('C1',)
 for i, step in enumerate(steps):