jinns 0.9.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

jinns/{data/_display.py → plot/_plot.py}

@@ -1,5 +1,5 @@
 """
-Utility functions for plotting
+Utility functions for plotting in 1D and 2D, with and without time.
 """
 
 from functools import partial
@@ -8,48 +8,55 @@ import matplotlib.pyplot as plt
 import jax.numpy as jnp
 from jax import vmap
 from mpl_toolkits.axes_grid1 import ImageGrid
+from typing import Callable, List
+from jaxtyping import Array, Float, Bool
 
 
 def plot2d(
-    fun,
-    xy_data,
-    times=None,
-    Tmax=1,
-    title="",
-    figsize=(7, 7),
-    cmap="inferno",
-    spinn=False,
-    vmin_vmax=None,
-    ax_for_plot=None,
+    fun: Callable,
+    xy_data: tuple[Float[Array, "nx"], Float[Array, "ny"]],
+    times: Float[Array, "nt"] | List[float] | None = None,
+    Tmax: float = 1,
+    title: str = "",
+    figsize: tuple = (7, 7),
+    cmap: str = "inferno",
+    spinn: bool = False,
+    vmin_vmax: tuple[float, float] | None = None,
+    ax_for_plot: plt.Axes | None = None,
 ):
     r"""Generic function for plotting functions over rectangular 2-D domains
-    :math:`\Omega`. It treats both the stationary case :math:`u(x)` or the
-    non-stationnary case :math:`u(t, x)`.
+    $\Omega$. It handles both the
 
-    When in the non-stationnary case, the `times` argument gives the time
-    slices :math:`t_i` at which to plot :math:`u(t_i, x)`.
+     1. the stationary case $u(x)$
+     2. the non-stationnary case $u(t, x)$
+
+    In the non-stationnary case, the `times` argument gives the time
+    slices $t_i$ at which to plot $u(t_i, x)$.
 
 
     Parameters
     ----------
-    fun : _type_
-        _description_
-    xy_data : _type_
-        _description_
-    times : _type_, optional
-        _description_, by default None
-    Tmax : float, only in non-stationary cases
-        Useful if you used time re-scaling in the differential equation, by
-        default 1
-    title : str, optional
+    fun :
+        the function $u$ to plot on the meshgrid, and eventually the time
+        slices. It's suppose to have signature `u(x)` in the stationnary case,, and `u(t, x)` in the non-stationnary case. Use `partial` or `lambda to freeze / reorder any other arguments.
+    xy_data :
+        A list of 2 `jnp.Array` providing grid values for meshgrid creation
+    times :
+        list or Array of time slices where to plot the function. Use Tmax if
+        you trained with time-rescaling.
+    Tmax :
+        Useful if you used time rescaling in the differential equation for training, default to 1 (no rescaling).
+    title :
         plot title, by default ""
-    figsize : tuple, optional
-        _description_, by default (7, 7)
-    cmap : str, optional
-        _description_, by default "inferno"
-    vmin_vmax : tuple, optional
+    figsize :
+        By default (7, 7)
+    cmap :
+        the matplotlib color map used in the ImageGrid.
+    vmin_vmax :
         The colorbar minimum and maximum value. Defaults None.
-    ax_for_plot : Matplotlib axis, optional
+    spinn :
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+    ax_for_plot :
         If None, jinns triggers the plotting. Otherwise this argument
         corresponds to the axis which will host the plot. Default is None.
         NOTE: that this argument will have an effect only if times is None.
@@ -57,7 +64,7 @@ def plot2d(
     Raises
     ------
     ValueError
-        if xy_data is not a list of type 2
+        if xy_data is not a list of length 2
     """
 
     # if not isinstance(xy_data, jnp.ndarray) and not xy_data.shape[-1] == 2:
@@ -179,13 +186,13 @@ def plot2d(
 
 def _plot_2D_statio(
     v_fun,
-    mesh,
-    plot=True,
-    colorbar=True,
-    cmap="inferno",
-    figsize=(7, 7),
-    spinn=False,
-    vmin_vmax=None,
+    mesh: Float[Array, "nx*ny nx*ny"],
+    plot: Bool = True,
+    colorbar: Bool = True,
+    cmap: str = "inferno",
+    figsize: tuple[int, int] = (7, 7),
+    spinn: Bool = False,
+    vmin_vmax: tuple[float, float] = None,
 ):
     """Function that plot the function u(x) with 2-D input x using pcolormesh()
 
@@ -200,6 +207,12 @@ def _plot_2D_statio(
         either show or return the plot, by default True
     colorbar : bool, optional
         add a colorbar, by default True
+    cmap :
+        the matplotlib color map used in the ImageGrid.
+    figsize :
+        By default (7, 7)
+    spinn :
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
     vmin_vmax: tuple, optional
         The colorbar minimum and maximum value. Defaults None.
 
@@ -238,33 +251,37 @@ def _plot_2D_statio(
 
 
 def plot1d_slice(
-    fun,
-    xdata,
-    time_slices=jnp.array([0]),
-    Tmax=1,
-    title="",
-    figsize=(10, 10),
-    spinn=False,
+    fun: Callable[[float, float], float],
+    xdata: Float[Array, "nx"],
+    time_slices: Float[Array, "nt"] | None = None,
+    Tmax: float = 1.0,
+    title: str = "",
+    figsize: tuple[int, int] = (10, 10),
+    spinn: Bool = False,
 ):
     """Function for plotting time slices of a function :math:`f(t_i, x)` where
-    `t` is time (1-D) and x is 1-D
+    `t_i` is time (1-D) and x is 1-D
 
     Parameters
     ----------
-    fun : callable with two arguments `t` and `x`
+    fun
         f(t, x)
-    xdata : jnp.array
+    xdata
         the discretization of space
-    time_slices : list, optional
-        the time slices :math:`t_i` at which to plot, by default [0]
-    Tmax : int, optional
+    time_slices
+        the time slices :math:`t_i` at which to plot.
+    Tmax
         Useful if you used time re-scaling in the differential equation, by
         default 1
-    title : str, optional
+    title
         title of the plot, by default ""
-    figsize : tuple, optional
+    spinn :
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+    figsize
         size of the figure, by default (10, 10)
     """
+    if time_slices is None:
+        time_slices = jnp.array([0])
     plt.figure(figsize=figsize)
     for t in time_slices:
         if not spinn:
@@ -284,16 +301,16 @@ def plot1d_slice(
 
 
 def plot1d_image(
-    fun,
-    xdata,
-    times,
-    Tmax=1,
-    title="",
-    figsize=(10, 10),
-    colorbar=True,
-    cmap="inferno",
-    spinn=False,
-    vmin_vmax=None,
+    fun: Callable[[float, float], float],
+    xdata: Float[Array, "nx"],
+    times: Float[Array, "nt"],
+    Tmax: float = 1.0,
+    title: str = "",
+    figsize: tuple[int, int] = (10, 10),
+    colorbar: Bool = True,
+    cmap: str = "inferno",
+    spinn: Bool = False,
+    vmin_vmax: tuple[float, float] = None,
 ):
     """Function for plotting the 2-D image of a function :math:`f(t, x)` where
     `t` is time (1-D) and x is space (1-D).
@@ -302,19 +319,25 @@ def plot1d_image(
 
     Parameters
     ----------
-    fun : callable with two arguments t and x
-        the function to plot
-    xdata : jnp.array
+    fun :
+        callable with two arguments t and x the function to plot
+    xdata :
         the discretization of space
-    times : jnp.array
+    times :
         the discretization of time
-    Tmax : int, optional
-        _description_, by default 1
-    title : str, optional
-        , by default ""
-    figsize : tuple, optional
-        , by default (10, 10)
-    vmin_vmax: tuple
+    Tmax :
+        by default 1
+    title :
+        by default ""
+    figsize :
+        by default (10, 10)
+    colorbar :
+        Whether to add a colobar
+    cmap :
+        the matplotlib color map used in the ImageGrid.
+    spinn :
+        True if a SPINN is to be plotted. False for PINNs and HYPERPINNs
+    vmin_vmax:
         The colorbar minimum and maximum value. Defaults None.
     """
 

jinns/solver/_rar.py

@@ -1,21 +1,32 @@
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+
+from typing import TYPE_CHECKING, Callable
+from functools import partial
 import jax
 from jax import vmap
 import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Int, Bool
+
+from jinns.data._Batchs import *
+from jinns.loss._LossODE import LossODE, SystemLossODE
+from jinns.loss._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
 from jinns.data._DataGenerators import (
     DataGeneratorODE,
     CubicMeshPDEStatio,
     CubicMeshPDENonStatio,
 )
-from jinns.loss._LossPDE import LossPDEStatio, LossPDENonStatio, SystemLossPDE
-from jinns.loss._LossODE import LossODE, SystemLossODE
-from jinns.loss._DynamicLossAbstract import PDEStatio
-
-from functools import partial
 from jinns.utils._hyperpinn import HYPERPINN
 from jinns.utils._spinn import SPINN
 
 
-def _proceed_to_rar(data, i):
+if TYPE_CHECKING:
+    from jinns.utils._types import *
+
+
+def _proceed_to_rar(data: AnyDataGenerator, i: Int) -> Bool:
     """Utilility function with various check to ensure we can proceed with the rar_step.
     Return True if yes, and False otherwise"""
 
@@ -30,13 +41,13 @@ def _proceed_to_rar(data, i):
     # Memory allocation checks (depends on the type of DataGenerator)
     # check if we still have room to append new collocation points in the
     # allocated jnp.array (can concern `data.p_times` or `p_omega`)
-    if isinstance(data, DataGeneratorODE) or isinstance(data, CubicMeshPDENonStatio):
+    if isinstance(data, (DataGeneratorODE, CubicMeshPDENonStatio)):
         check_list.append(
             data.rar_parameters["selected_sample_size_times"]
             <= jnp.count_nonzero(data.p_times == 0),
         )
 
-    if isinstance(data, CubicMeshPDEStatio) or isinstance(data, CubicMeshPDENonStatio):
+    if isinstance(data, (CubicMeshPDEStatio, CubicMeshPDENonStatio)):
         # for now the above check are redundants but there may be a time when
         # we drop inheritence
         check_list.append(
@@ -49,7 +60,14 @@ def _proceed_to_rar(data, i):
 
 
 @partial(jax.jit, static_argnames=["_rar_step_true", "_rar_step_false"])
-def trigger_rar(i, loss, params, data, _rar_step_true, _rar_step_false):
+def trigger_rar(
+    i: Int,
+    loss: AnyLoss,
+    params: AnyParams,
+    data: AnyDataGenerator,
+    _rar_step_true: Callable[[rar_operands], AnyDataGenerator],
+    _rar_step_false: Callable[[rar_operands], AnyDataGenerator],
+) -> tuple[AnyLoss, AnyParams, AnyDataGenerator]:
 
     if data.rar_parameters is None:
         # do nothing.
@@ -65,7 +83,13 @@ def trigger_rar(i, loss, params, data, _rar_step_true, _rar_step_false):
         return loss, params, data
 
 
-def init_rar(data):
+def init_rar(
+    data: AnyDataGenerator,
+) -> tuple[
+    AnyDataGenerator,
+    Callable[[rar_operands], AnyDataGenerator],
+    Callable[[rar_operands], AnyDataGenerator],
+]:
     """
     Separated from the main rar, because the initialization to get _true and
     _false cannot be jit-ted.
@@ -100,13 +124,21 @@ def init_rar(data):
                 data.rar_parameters["sample_size_omega"],
                 data.rar_parameters["selected_sample_size_omega"],
             )
+        else:
+            raise ValueError(f"Wrong type for data got {type(data)}")
 
-        data.rar_parameters["iter_from_last_sampling"] = 0
+        if isinstance(data, eqx.Module):
+            data = eqx.tree_at(lambda m: m.rar_iter_from_last_sampling, data, 0)
+        else:
+            data.rar_iter_from_last_sampling = 0
 
     return data, _rar_step_true, _rar_step_false
 
 
-def _rar_step_init(sample_size, selected_sample_size):
+def _rar_step_init(sample_size: Int, selected_sample_size: Int) -> tuple[
+    Callable[[rar_operands], AnyDataGenerator],
+    Callable[[rar_operands], AnyDataGenerator],
+]:
     """
     This is a wrapper because the sampling size and
     selected_sample_size, must be treated as static
@@ -116,11 +148,17 @@ def _rar_step_init(sample_size, selected_sample_size):
     This is a kind of manual declaration of static argnums
     """
 
-    def rar_step_true(operands):
+    def rar_step_true(operands: rar_operands) -> AnyDataGenerator:
         loss, params, data, i = operands
 
         if isinstance(data, DataGeneratorODE):
-            new_omega_samples = data.sample_in_time_domain(sample_size)
+
+            if isinstance(data, eqx.Module):
+                new_key, subkey = jax.random.split(data.key)
+                new_omega_samples = data.sample_in_time_domain(subkey, sample_size)
+                data = eqx.tree_at(lambda m: m.key, data, new_key)
+            else:
+                new_omega_samples = data.sample_in_time_domain(sample_size)
 
             # We can have different types of Loss
             if isinstance(loss, LossODE):
@@ -162,17 +200,29 @@ def _rar_step_init(sample_size, selected_sample_size):
             ## add the new points in times
             # start indices of update can be dynamic but the the shape (length)
             # of the slice
-            data.times = jax.lax.dynamic_update_slice(
+            new_times = jax.lax.dynamic_update_slice(
                 data.times,
                 higher_residual_points,
                 (data.nt_start + data.rar_iter_nb * selected_sample_size,),
             )
 
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(lambda m: m.times, data, new_times)
+            else:
+                data.times = new_times
             ## rearrange probabilities so that the probabilities of the new
             ## points are non-zero
             new_proba = 1 / (data.nt_start + data.rar_iter_nb * selected_sample_size)
             # the next work because nt_start is static
-            data.p_times = data.p_times.at[: data.nt_start].set(new_proba)
+            new_p_times = data.p_times.at[: data.nt_start].set(new_proba)
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(
+                    lambda m: m.p_times,
+                    data,
+                    new_p_times,
+                )
+            else:
+                data.p_times = new_p_times
 
             # the next requires a fori_loop because the range is dynamic
             def update_slices(i, p):
@@ -182,16 +232,29 @@ def _rar_step_init(sample_size, selected_sample_size):
                     ((data.nt_start + i * selected_sample_size),),
                 )
 
-            data.rar_iter_nb += 1
-
-            data.p_times = jax.lax.fori_loop(
+            new_rar_iter_nb = data.rar_iter_nb + 1
+            new_p_times = jax.lax.fori_loop(
                 0, data.rar_iter_nb, update_slices, data.p_times
             )
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(
+                    lambda m: (m.rar_iter_nb, m.p_times),
+                    data,
+                    (new_rar_iter_nb, new_p_times),
+                )
+            else:
+                data.rar_iter_nb = new_rar_iter_nb
+                data.p_times = new_p_times
 
         elif isinstance(data, CubicMeshPDEStatio) and not isinstance(
             data, CubicMeshPDENonStatio
         ):
-            new_omega_samples = data.sample_in_omega_domain(sample_size)
+            if isinstance(data, eqx.Module):
+                new_key, *subkeys = jax.random.split(data.key, data.dim + 1)
+                new_omega_samples = data.sample_in_omega_domain(subkeys, sample_size)
+                data = eqx.tree_at(lambda m: m.key, data, new_key)
+            else:
+                new_omega_samples = data.sample_in_omega_domain(sample_size)
 
             # We can have different types of Loss
             if isinstance(loss, LossPDEStatio):
@@ -209,7 +272,7 @@ def _rar_step_init(sample_size, selected_sample_size):
                     mse_on_s = (jnp.linalg.norm(dyn_on_s, axis=-1) ** 2).flatten()
                 else:
                     mse_on_s = dyn_on_s**2
-            elif isinstance(loss, SystemLossPDE):
+            elif isinstance(loss, SystemLossODE):
                 mse_on_s = 0
                 for i in loss.dynamic_loss_dict.keys():
                     # only the case LossPDEStatio here
@@ -237,17 +300,30 @@ def _rar_step_init(sample_size, selected_sample_size):
             ## add the new points in omega
             # start indices of update can be dynamic but not the shape (length)
             # of the slice
-            data.omega = jax.lax.dynamic_update_slice(
+            new_omega = jax.lax.dynamic_update_slice(
                 data.omega,
                 higher_residual_points,
                 (data.n_start + data.rar_iter_nb * selected_sample_size, data.dim),
             )
 
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(lambda m: m.omega, data, new_omega)
+            else:
+                data.omega = new_omega
+
             ## rearrange probabilities so that the probabilities of the new
             ## points are non-zero
             new_proba = 1 / (data.n_start + data.rar_iter_nb * selected_sample_size)
             # the next work because n_start is static
-            data.p_omega = data.p_omega.at[: data.n_start].set(new_proba)
+            new_p_omega = data.p_omega.at[: data.n_start].set(new_proba)
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(
+                    lambda m: m.p_omega,
+                    data,
+                    new_p_omega,
+                )
+            else:
+                data.p_omega = new_p_omega
 
             # the next requires a fori_loop because the range is dynamic
             def update_slices(i, p):
@@ -257,11 +333,19 @@ def _rar_step_init(sample_size, selected_sample_size):
                     ((data.n_start + i * selected_sample_size),),
                 )
 
-            data.rar_iter_nb += 1
-
-            data.p_omega = jax.lax.fori_loop(
+            new_rar_iter_nb = data.rar_iter_nb + 1
+            new_p_omega = jax.lax.fori_loop(
                 0, data.rar_iter_nb, update_slices, data.p_omega
             )
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(
+                    lambda m: (m.rar_iter_nb, m.p_omega),
+                    data,
+                    (new_rar_iter_nb, new_p_omega),
+                )
+            else:
+                data.rar_iter_nb = new_rar_iter_nb
+                data.p_omega = new_p_omega
 
         elif isinstance(data, CubicMeshPDENonStatio):
             if isinstance(loss.u, HYPERPINN) or isinstance(loss.u, SPINN):
@@ -274,8 +358,19 @@ def _rar_step_init(sample_size, selected_sample_size):
             )
             sample_size_times, sample_size_omega = sample_size
 
-            new_times_samples = data.sample_in_time_domain(sample_size_times)
-            new_omega_samples = data.sample_in_omega_domain(sample_size_omega)
+            if isinstance(data, eqx.Module):
+                new_key, subkey = jax.random.split(data.key)
+                new_times_samples = data.sample_in_time_domain(
+                    subkey, sample_size_times
+                )
+                new_key, *subkeys = jax.random.split(new_key, data.dim + 1)
+                new_omega_samples = data.sample_in_omega_domain(
+                    subkeys, sample_size_omega
+                )
+                data = eqx.tree_at(lambda m: m.key, data, new_key)
+            else:
+                new_times_samples = data.sample_in_time_domain(sample_size_times)
+                new_omega_samples = data.sample_in_omega_domain(sample_size_omega)
 
             if not data.cartesian_product:
                 times = new_times_samples
@@ -342,14 +437,23 @@ def _rar_step_init(sample_size, selected_sample_size):
             ## add the new points in times
             # start indices of update can be dynamic but not the shape (length)
             # of the slice
-            data.times = jax.lax.dynamic_update_slice(
+            new_times = jax.lax.dynamic_update_slice(
                 data.times,
                 higher_residual_points_times,
-                (data.n_start + data.rar_iter_nb * selected_sample_size_times,),
+                (
+                    data.n_start
+                    + data.rar_iter_nb  # NOTE typo here nt_start ?
+                    * selected_sample_size_times,
+                ),
             )
 
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(lambda m: m.times, data, new_times)
+            else:
+                data.times = new_times
+
             ## add the new points in omega
-            data.omega = jax.lax.dynamic_update_slice(
+            new_omega = jax.lax.dynamic_update_slice(
                 data.omega,
                 higher_residual_points_omega,
                 (
@@ -358,19 +462,38 @@ def _rar_step_init(sample_size, selected_sample_size):
                 ),
             )
 
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(lambda m: m.omega, data, new_omega)
+            else:
+                data.omega = new_omega
+
             ## rearrange probabilities so that the probabilities of the new
             ## points are non-zero
             new_p_times = 1 / (
                 data.nt_start + data.rar_iter_nb * selected_sample_size_times
             )
             # the next work because nt_start is static
-            data.p_times = data.p_times.at[: data.nt_start].set(new_p_times)
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(
+                    lambda m: m.p_times,
+                    data,
+                    data.p_times.at[: data.nt_start].set(new_p_times),
+                )
+            else:
+                data.p_times = data.p_times.at[: data.nt_start].set(new_p_times)
 
             # same for p_omega (work because n_start is static)
             new_p_omega = 1 / (
                 data.n_start + data.rar_iter_nb * selected_sample_size_omega
             )
-            data.p_omega = data.p_omega.at[: data.n_start].set(new_p_omega)
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(
+                    lambda m: m.p_omega,
+                    data,
+                    data.p_omega.at[: data.n_start].set(new_p_omega),
+                )
+            else:
+                data.p_omega = data.p_omega.at[: data.n_start].set(new_p_omega)
 
             # the part of data.p_* after n_start requires a fori_loop because
             # the range is dynamic
@@ -385,13 +508,13 @@ def _rar_step_init(sample_size, selected_sample_size):
 
                 return update_slices
 
-            data.rar_iter_nb += 1
+            new_rar_iter_nb = data.rar_iter_nb + 1
 
             ## update rest of p_times
             update_slices_times = create_update_slices(
                 new_p_times, selected_sample_size_times
             )
-            data.p_times = jax.lax.fori_loop(
+            new_p_times = jax.lax.fori_loop(
                 0,
                 data.rar_iter_nb,
                 update_slices_times,
@@ -401,21 +524,34 @@ def _rar_step_init(sample_size, selected_sample_size):
             update_slices_omega = create_update_slices(
                 new_p_omega, selected_sample_size_omega
             )
-            data.p_omega = jax.lax.fori_loop(
+            new_p_omega = jax.lax.fori_loop(
                 0,
                 data.rar_iter_nb,
                 update_slices_omega,
                 data.p_omega,
             )
+            if isinstance(data, eqx.Module):
+                data = eqx.tree_at(
+                    lambda m: (m.rar_iter_nb, m.p_omega, m.p_times),
+                    data,
+                    (new_rar_iter_nb, new_p_omega, new_p_times),
+                )
+            else:
+                data.rar_iter_nb = new_rar_iter_nb
+                data.p_times = new_p_times
+                data.p_omega = new_p_omega
 
         # update RAR parameters for all cases
-        data.rar_iter_from_last_sampling = 0
+        if isinstance(data, eqx.Module):
+            data = eqx.tree_at(lambda m: m.rar_iter_from_last_sampling, data, 0)
+        else:
+            data.rar_iter_from_last_sampling = 0
 
         # NOTE must return data to be correctly updated because we cannot
         # have side effects in this function that will be jitted
         return data
 
-    def rar_step_false(operands):
+    def rar_step_false(operands: rar_operands) -> AnyDataGenerator:
         _, _, data, i = operands
 
         # Add 1 only if we are after the burn in period
@@ -425,7 +561,15 @@ def _rar_step_init(sample_size, selected_sample_size):
             lambda: 1,
         )
 
-        data.rar_iter_from_last_sampling += increment
+        new_rar_iter_from_last_sampling = data.rar_iter_from_last_sampling + increment
+        if isinstance(data, eqx.Module):
+            data = eqx.tree_at(
+                lambda m: m.rar_iter_from_last_sampling,
+                data,
+                new_rar_iter_from_last_sampling,
+            )
+        else:
+            data.rar_iter_from_last_sampling = new_rar_iter_from_last_sampling
         return data
 
     return rar_step_true, rar_step_false