Skip to content

generate_eq_sols

Functions:

  • eq_sol –

    Get the solution of the equilibrium equations for given parameters.

  • generate_equilibrium_solutions –

    Generate equilibrium solution objects for a set of parameters and save them to files.

  • main –

    Command-line interface for generating and saving equilibrium solutions.

eq_sol 

eq_sol(alpha: float, theta_0: float, delta_h: float) -> Tuple[OdeSolution, float, bool]

Get the solution of the equilibrium equations for given parameters.

Parameters:

  • alpha

    (float) –

    Alpha parameter.

  • theta_0

    (float) –

    Theta_0 parameter.

  • delta_h

    (float) –

    Delta_h parameter.

Returns:

  • sol ( OdeSolution ) –

    Solution object of the second equilibrium ODE system.

  • delta_a ( float ) –

    Computed value of the delta_a parameter obtained from the first equilibrium solution.

  • rfp_flag ( bool ) –

    Flag indicating whether the resulting configuration is in the RFP (reversed field pinch) regime according to :func:eqs.eq2_sol.

Source code in src/fpga_profile_reco/data/generate_eq_sols.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def eq_sol(alpha: float, theta_0: float, delta_h: float) -> Tuple[OdeSolution, float, bool]:
    """
    Get the solution of the equilibrium equations for given parameters.

    Parameters
    ----------
    alpha : float
        Alpha parameter.
    theta_0 : float
        Theta_0 parameter.
    delta_h : float
        Delta_h parameter.

    Returns
    -------
    sol : scipy.integrate.OdeSolution
        Solution object of the second equilibrium ODE system.
    delta_a : float
        Computed value of the ``delta_a`` parameter obtained from the first
        equilibrium solution.
    rfp_flag : bool
        Flag indicating whether the resulting configuration is in the RFP
        (reversed field pinch) regime according to :func:`eqs.eq2_sol`.
    """
    sol_1 = eqs.eq1_sol(alpha, theta_0, delta_h)
    delta_a = sol_1(eqs.RA - np.abs(delta_h))[2]
    sol, rfp_flag = eqs.eq2_sol(alpha, theta_0, delta_h, delta_a)

    return sol, delta_a, rfp_flag

generate_equilibrium_solutions 

generate_equilibrium_solutions(n_samples: int, grid: bool, alpha_samples: int, theta_0_samples: int, delta_h_samples: int, seed: int, save_path: Path, n_jobs: int = 1)

Generate equilibrium solution objects for a set of parameters and save them to files.

Depending on the value of grid, this function either generates a regular grid over the parameter space or draws random samples uniformly within the parameter bounds defined in :mod:fpga_profile_reco.data.equations. For each parameter combination, the corresponding equilibrium ODE solutions are computed and serialized in batched pickle files under save_path.

Parameters:

  • n_samples

    (int) –

    Total number of random samples to generate (only used if grid is False).

  • grid

    (bool) –

    If True, generate a Cartesian grid of parameters using alpha_samples, theta_0_samples and delta_h_samples. Otherwise, sample parameters randomly.

  • alpha_samples

    (int) –

    Number of samples for the alpha parameter when grid is True.

  • theta_0_samples

    (int) –

    Number of samples for the theta_0 parameter when grid is True.

  • delta_h_samples

    (int) –

    Number of samples for the delta_h parameter when grid is True.

  • seed

    (int) –

    Random seed for reproducibility (only used if grid is False).

  • save_path

    (Path) –

    Directory where the generated solution batches (pickle files) are saved. The directory is created if it does not exist.

  • n_jobs

    (int, default: 1 ) –

    Number of parallel worker processes used to compute the solutions. Defaults to 1.

Returns:

  • None –

    The function is executed for its side effects of computing and saving solution batches to disk.

Source code in src/fpga_profile_reco/data/generate_eq_sols.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def generate_equilibrium_solutions(n_samples: int, grid: bool, alpha_samples: int, theta_0_samples: int, delta_h_samples: int, seed: int, save_path: Path, n_jobs: int = 1):
    """
    Generate equilibrium solution objects for a set of parameters and save them to files.

    Depending on the value of ``grid``, this function either generates a regular
    grid over the parameter space or draws random samples uniformly within the
    parameter bounds defined in :mod:`fpga_profile_reco.data.equations`. For each
    parameter combination, the corresponding equilibrium ODE solutions are
    computed and serialized in batched pickle files under ``save_path``.

    Parameters
    ----------
    n_samples : int
        Total number of random samples to generate (only used if ``grid`` is
        False).
    grid : bool
        If True, generate a Cartesian grid of parameters using
        ``alpha_samples``, ``theta_0_samples`` and ``delta_h_samples``.
        Otherwise, sample parameters randomly.
    alpha_samples : int
        Number of samples for the ``alpha`` parameter when ``grid`` is True.
    theta_0_samples : int
        Number of samples for the ``theta_0`` parameter when ``grid`` is True.
    delta_h_samples : int
        Number of samples for the ``delta_h`` parameter when ``grid`` is True.
    seed : int
        Random seed for reproducibility (only used if ``grid`` is False).
    save_path : pathlib.Path
        Directory where the generated solution batches (pickle files) are saved.
        The directory is created if it does not exist.
    n_jobs : int, optional
        Number of parallel worker processes used to compute the solutions.
        Defaults to 1.

    Returns
    -------
    None
        The function is executed for its side effects of computing and saving
        solution batches to disk.
    """

    if not save_path.exists():
        save_path.mkdir(parents=True, exist_ok=True)

    if grid:
        alpha = np.linspace(eqs.ALPHA_MIN, eqs.ALPHA_MAX, alpha_samples)
        theta_0 = np.linspace(eqs.THETA_0_MIN, eqs.THETA_0_MAX, theta_0_samples)
        delta_h = np.linspace(eqs.DELTA_H_MIN, eqs.DELTA_H_MAX, delta_h_samples)
        # create array of parameter combinations
        params = np.array(np.meshgrid(alpha, theta_0, delta_h)).T.reshape(-1, 3)
    else:
        # set seed for reproducibility
        np.random.seed(seed)
        # vectorized generation n_samples random parameter combinations
        params = np.random.uniform([eqs.ALPHA_MIN, eqs.THETA_0_MIN, eqs.DELTA_H_MIN], [eqs.ALPHA_MAX, eqs.THETA_0_MAX, eqs.DELTA_H_MAX], (n_samples, 3))

    batch_size = min(1000, len(params))
    batches = [params[i:i + batch_size] for i in range(0, len(params), batch_size)]

    with ProcessPoolExecutor(max_workers=n_jobs) as executor:
        futures = {executor.submit(generate_sol_batch, (i, save_path, batch)): i for i, batch in enumerate(batches)}
        for future in tqdm(as_completed(futures), total=len(futures), desc="Computing equilibrium solutions..."):
            batch_id = futures[future]
            try:
                future.result()  # Wait for the result and handle exceptions if any
            except Exception as e:
                print(f"Error processing batch {batch_id}: {e}")

main 

main()

Command-line interface for generating and saving equilibrium solutions.

This function parses command-line arguments and forwards them to :func:generate_equilibrium_solutions. It supports both grid-based and random sampling of parameters and controls the output directory and parallelism.

Command Line Parameters
  • --n_samples : int, optional Total number of random samples to generate (only used if --grid is not set). Default is 100000.
  • --grid : flag, optional If provided, generate a Cartesian grid of parameters instead of random samples.
  • --alpha_samples : int, optional Number of samples for the alpha parameter when --grid is set. Default is 101.
  • --theta_0_samples : int, optional Number of samples for the theta_0 parameter when --grid is set. Default is 101.
  • --delta_h_samples : int, optional Number of samples for the delta_h parameter when --grid is set. Default is 11.
  • --seed : int, optional Random seed for reproducibility when --grid is not set. Default is 42.
  • --save_path : pathlib.Path Path to the output directory where solution batches will be stored. This argument is required.
  • --n_jobs : int, optional Number of parallel worker processes. Default is 1.
Source code in src/fpga_profile_reco/data/generate_eq_sols.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def main():
    """
    Command-line interface for generating and saving equilibrium solutions.

    This function parses command-line arguments and forwards them to
    :func:`generate_equilibrium_solutions`. It supports both grid-based and
    random sampling of parameters and controls the output directory and
    parallelism.

    Command Line Parameters
    -----------------------
    - `--n_samples` : int, optional
        Total number of random samples to generate (only used if ``--grid`` is
        not set). Default is 100000.
    - `--grid` : flag, optional
        If provided, generate a Cartesian grid of parameters instead of random
        samples.
    - `--alpha_samples` : int, optional
        Number of samples for the ``alpha`` parameter when ``--grid`` is set.
        Default is 101.
    - `--theta_0_samples` : int, optional
        Number of samples for the ``theta_0`` parameter when ``--grid`` is set.
        Default is 101.
    - `--delta_h_samples` : int, optional
        Number of samples for the ``delta_h`` parameter when ``--grid`` is set.
        Default is 11.
    - `--seed` : int, optional
        Random seed for reproducibility when ``--grid`` is not set. Default is
        42.
    - `--save_path` : pathlib.Path
        Path to the output directory where solution batches will be stored.
        This argument is required.
    - `--n_jobs` : int, optional
        Number of parallel worker processes. Default is 1.
    """
    import argparse

    parser = argparse.ArgumentParser(description="Generate equilibrium solutions.")
    parser.add_argument("--n_samples", type=int, default=100_000, help="Total number of random samples to generate (only used if --grid is not set).")
    parser.add_argument("--grid", action='store_true', help="Generate a grid of parameters.")
    parser.add_argument("--alpha_samples", type=int, default=101, help="Number of samples for alpha parameter (for --grid option).")
    parser.add_argument("--theta_0_samples", type=int, default=101, help="Number of samples for theta_0 parameter (for --grid option).")
    parser.add_argument("--delta_h_samples", type=int, default=11, help="Number of samples for delta_h parameter (for --grid option).")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility. (only used if --grid is not set)")
    parser.add_argument("--save_path", type=Path, required=True, help="Path to save the generated solutions.")
    parser.add_argument("--n_jobs", type=int, default=1, help="Number of parallel jobs to run.")

    args = parser.parse_args()

    generate_equilibrium_solutions(args.n_samples, args.grid, args.alpha_samples, args.theta_0_samples, args.delta_h_samples, args.seed, args.save_path, args.n_jobs)