def generate_batch_dataset(
    batch_size,
):
    """Generate a batch of synthetic data with numerical safeguards."""
    I_batch = []
    dos_s_batch = []
    z_batch = []
    for i in range(batch_size):
        I_V, dos_s, z = single_data_generator()
        I_batch.append(I_V)
        dos_s_batch.append(dos_s)
        z_batch.append(z)

    return (
        tf.constant(I_V, dtype=tf.float32),
        tf.constant(dos_s, dtype=tf.float32),
        tf.constant(z, dtype=tf.float32),
    )


def generate_multiple_datasets(
    num_datasets=5,
    N_per_dataset=1000,
    batch_size=100,  # Process N_per_dataset samples in batches of batch_size
    num_points=100,
    V_max=10,
    z_means=[7.5, 8.0, 6.5],
    z_stds=[1.0, 0.8, 1.2],
    noise_levels=[0.01, 0.02, 0.01],
    rho_s_types=["gaussian", "step", "random"],
    save_format="tfrecord",
    output_dir="datasets",
):
    """
    Generate multiple synthetic datasets in batches with numerical safeguards.

    Args:
        num_datasets: Number of datasets to generate.
        N_per_dataset: Number of samples per dataset.
        batch_size: Number of samples processed at once (accelerates CPU operations).
        num_points: Number of points in E and V grids.
        V_max: Maximum bias voltage (eV).
        z_means: List of z means for each dataset.
        z_stds: List of z stds for each dataset.
        noise_levels: List of noise levels for each dataset.
        rho_s_types: List of rho_s types for each dataset.
        save_format: Format to save datasets ('tfrecord', 'numpy', or 'hdf5').
        output_dir: Directory to save datasets.
    """
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)

    def generate_and_save(i):
        """Helper function to generate and save a single dataset in batches."""
        # Initialize arrays for the entire dataset
        I_full = np.zeros((N_per_dataset, num_points))
        rho_s_full = np.zeros((N_per_dataset, num_points))
        z_full = np.zeros(N_per_dataset)

        # Process in batches
        for batch_start in range(0, N_per_dataset, batch_size):
            batch_end = min(batch_start + batch_size, N_per_dataset)
            batch_N = batch_end - batch_start

            # Generate batch
            I_batch, rho_s_batch, z_batch = generate_batch_dataset(
                N=batch_N,
                num_points=num_points,
                V_max=V_max,
                z_mean=z_means[i % len(z_means)],
                z_std=z_stds[i % len(z_stds)],
                noise_level=noise_levels[i % len(noise_levels)],
                rho_s_type=rho_s_types[i % len(rho_s_types)],
            )

            # Store batch results
            I_full[batch_start:batch_end] = I_batch.numpy()
            rho_s_full[batch_start:batch_end] = rho_s_batch.numpy()
            z_full[batch_start:batch_end] = z_batch.numpy()

        # Normalize and clip
        I_full = I_full / np.max(I_full, axis=1, keepdims=True)
        rho_s_full = rho_s_full / np.max(rho_s_full, axis=1, keepdims=True)
        z_full = np.clip(z_full, 5.0, 10.0)

        # Save based on format
        if save_format == "tfrecord":
            filename = os.path.join(output_dir, f"dataset_{i}.tfrecord")
            with tf.io.TFRecordWriter(filename) as writer:
                for j in range(N_per_dataset):
                    example = tf.train.Example(
                        features=tf.train.Features(
                            feature={
                                "I": tf.train.Feature(
                                    float_list=tf.train.FloatList(value=I_full[j])
                                ),
                                "rho_s": tf.train.Feature(
                                    float_list=tf.train.FloatList(value=rho_s_full[j])
                                ),
                                "z": tf.train.Feature(
                                    float_list=tf.train.FloatList(value=[z_full[j]])
                                ),
                            }
                        )
                    ).SerializeToString()
                    writer.write(example)
        elif save_format == "numpy":
            np.save(os.path.join(output_dir, f"I_{i}.npy"), I_full)
            np.save(os.path.join(output_dir, f"rho_s_{i}.npy"), rho_s_full)
            np.save(os.path.join(output_dir, f"z_{i}.npy"), z_full)
        elif save_format == "hdf5":
            import h5py

            filename = os.path.join(output_dir, f"dataset_{i}.h5")
            with h5py.File(filename, "w") as f:
                f.create_dataset("I", data=I_full)
                f.create_dataset("rho_s", data=rho_s_full)
                f.create_dataset("z", data=z_full)

    # Generate datasets in parallel
    with ThreadPoolExecutor() as executor:
        executor.map(generate_and_save, range(num_datasets))


the function single_data_generator() gives 2 arrays : I dos and 1 parameter z

refine the code