feat: store r_out as per-sample output when varied by optimization

    variable_param_dict = {key: param_dict[key] for key in variable_param_names}

    locked_param_dict = {key: param_dict[key] for key in locked_param_names}

    # Check if r_out varies across samples (e.g. due to optimization)

    # by computing the std of the last pipe value across all samples

    r_out_varies = False

    if "r_out" in locked_param_dict:

        r_out_last_values = np.array(

            [np.asarray(locked_param_dict["r_out"][s])[-1] for s in range(n_samples)]

        )

        if round(float(np.std(r_out_last_values)), 7) != 0:

            r_out_varies = True

    # Variable parameters: Create dataset with [param, samples] dimensions

    variable_param_da = xr.DataArray(

        list(variable_param_dict.values()),

    }

    pipe_pos = np.asarray(locked_param_x["pos"])

    r_out = locked_param_x["r_out"]

    pipe_pos_ds = xr.Dataset(

        data_vars=dict(pipe_pos=(["nPipes", "xy"], pipe_pos)),

        coords=dict(nPipes=(["nPipes"], pipes_coord)),

    )

    # Store r_out: if it varies across samples (optimization), store with [samples, nPipes] dims

    # otherwise store as locked with [nPipes] dim only

    if r_out_varies:

        r_out_per_sample = np.array(

            [np.asarray(locked_param_dict["r_out"][s]) for s in range(n_samples)]

        )

        r_out_ds = xr.Dataset(

            data_vars=dict(r_out=(["samples", "nPipes"], r_out_per_sample)),

            coords=dict(samples=(["samples"], list(sample_coord)), nPipes=(["nPipes"], pipes_coord)),

        )

    else:

        r_out = locked_param_x.get("r_out", locked_param_single_dict.get("r_out"))

        if r_out is not None and isinstance(r_out, list):

            r_out_ds = xr.Dataset(

                data_vars=dict(r_out=(["nPipes"], r_out)),

                coords=dict(nPipes=(["nPipes"], pipes_coord)),

            )

        else:

            r_out_ds = None

    # Then, create dataset from locked parameters with dim(len(parnames), nlayers_k_s)

    locked_param_z_k_s = {