bem_injection_random.f90 Source File

!> 粒子注入で共有する乱数初期化と基本分布sampling。
module bem_injection_random
  use bem_kinds, only: dp, i32, i64
  use bem_constants, only: k_boltzmann
  use bem_counter_rng, only: counter_stream_type
  use bem_random_stream, only: seed_random_stream, random_uniform
  implicit none
  private

  real(dp), parameter :: default_velocity_sigma_cutoff = 6.0_dp

  public :: seed_rng
  public :: sample_uniform_positions
  public :: sample_shifted_maxwell_velocities
  public :: sample_standard_normal

contains

  !> 与えたシード列またはシステム時刻から逐次 PCG64 乱数列を初期化する。
  !! @param[in] seed 乱数生成器へ与えるシード列（省略時は `system_clock` 値から生成）。
  subroutine seed_rng(seed)
    integer(i32), intent(in), optional :: seed(:)
    integer(i64) :: clk

    if (present(seed)) then
      call seed_random_stream(seed)
    else
      call system_clock(count=clk)
      call seed_random_stream([int(iand(clk, 2147483647_i64), i32), int(iand(ishft(clk, -31), 2147483647_i64), i32)])
    end if
  end subroutine seed_rng

  !> 直方体領域 `[low, high]` 内で一様分布の初期位置をサンプリングする。
  !! @param[in] low サンプリング領域の下限座標 `(x,y,z)` [m]。
  !! @param[in] high サンプリング領域の上限座標 `(x,y,z)` [m]。
  !! @param[out] x 生成した粒子位置配列 `x(3,n)` [m]。
  subroutine sample_uniform_positions(low, high, x)
    real(dp), intent(in) :: low(3), high(3)
    real(dp), intent(out) :: x(:, :)
    real(dp) :: span(3)
    integer :: i

    if (size(x, 1) /= 3) error stop "x first dimension must be 3"
    if (any(high < low)) error stop "high must be >= low for all axes"

    span = high - low
    call random_uniform(x)
    do i = 1, size(x, 2)
      x(:, i) = low + span*x(:, i)
    end do
  end subroutine sample_uniform_positions

  !> ドリフト速度付きMaxwell分布(温度または熱速度指定)から粒子速度を生成する。
  !! @param[in] drift_velocity 付加する平均ドリフト速度ベクトル `(vx,vy,vz)` [m/s]。
  !! @param[in] m_particle 粒子1個あたりの質量 [kg]（`temperature_k` 指定時に使用）。
  !! @param[out] v サンプリングした速度配列 `v(3,n)` [m/s]。
  !! @param[in] temperature_k 熱運動の温度 [K]（`thermal_speed` 未指定時に使用）。
  !! @param[in] thermal_speed 熱速度の標準偏差 `sigma` [m/s]（指定時は温度より優先）。
  !! @param[in] sigma_cutoff 標準正規変量を `[-sigma_cutoff, sigma_cutoff]` に切る上限（省略時 6）。
  subroutine sample_shifted_maxwell_velocities(drift_velocity, m_particle, v, temperature_k, thermal_speed, sigma_cutoff)
    real(dp), intent(in) :: drift_velocity(3)
    real(dp), intent(in) :: m_particle
    real(dp), intent(out) :: v(:, :)
    real(dp), intent(in), optional :: temperature_k
    real(dp), intent(in), optional :: thermal_speed
    real(dp), intent(in), optional :: sigma_cutoff
    integer :: i
    real(dp) :: sigma, cutoff

    if (size(v, 1) /= 3) error stop "v first dimension must be 3"
    if (m_particle <= 0.0_dp) error stop "particle mass [kg] must be > 0"

    if (.not. present(temperature_k) .and. .not. present(thermal_speed)) then
      error stop "either temperature [K] or thermal_speed must be provided"
    end if

    if (present(thermal_speed)) then
      if (thermal_speed < 0.0_dp) error stop "thermal_speed must be >= 0"
      sigma = thermal_speed
    else
      if (temperature_k < 0.0_dp) error stop "temperature [K] must be >= 0"
      sigma = sqrt(k_boltzmann*temperature_k/m_particle)
    end if
    cutoff = default_velocity_sigma_cutoff
    if (present(sigma_cutoff)) cutoff = sigma_cutoff

    ! Zero temperature still consumes the same draws for restart and trial replay.
    call sample_standard_normal(v, sigma_cutoff=cutoff)
    do i = 1, size(v, 2)
      v(:, i) = sigma*v(:, i) + drift_velocity
    end do
  end subroutine sample_shifted_maxwell_velocities

  !> Box–Muller法で標準正規乱数を生成し、任意形状配列へ詰める。
  !! @param[out] z 平均0・分散1の標準正規乱数で埋める出力配列。
  subroutine sample_standard_normal(z, sigma_cutoff, rng_seed, rng_counter)
    real(dp), intent(out) :: z(:, :)
    real(dp), intent(in), optional :: sigma_cutoff
    integer(i32), intent(in), optional :: rng_seed
    integer(i64), intent(in), optional :: rng_counter(:)
    type(counter_stream_type) :: stream
    integer :: n_total, i, row, column
    real(dp) :: r, theta, pi, u1, u2, z1, z2, cutoff

    ! Counter draws 100, 101, ... feed the Box-Muller pairs in order.
    if (present(rng_counter)) call stream%open(rng_seed, rng_counter, 100_i64)
    n_total = size(z)
    pi = acos(-1.0_dp)
    cutoff = default_velocity_sigma_cutoff
    if (present(sigma_cutoff)) cutoff = sigma_cutoff
    if (cutoff <= 0.0_dp) error stop "sigma_cutoff must be > 0"

    i = 1
    row = 1
    column = 1
    do while (i <= n_total)
      if (present(rng_counter)) then
        u1 = stream%next_uniform()
        u2 = stream%next_uniform()
      else
        call random_uniform(u1)
        call random_uniform(u2)
      end if
      if (u1 <= tiny(1.0_dp)) u1 = tiny(1.0_dp)
      r = sqrt(-2.0_dp*log(u1))
      theta = 2.0_dp*pi*u2
      z1 = r*cos(theta)
      z2 = r*sin(theta)
      if (abs(z1) <= cutoff) then
        z(row, column) = z1
        row = row + 1
        if (row > size(z, 1)) then
          row = 1
          column = column + 1
        end if
        i = i + 1
      end if
      if (i <= n_total .and. abs(z2) <= cutoff) then
        z(row, column) = z2
        row = row + 1
        if (row > size(z, 1)) then
          row = 1
          column = column + 1
        end if
        i = i + 1
      end if
    end do
    call stream%close()
  end subroutine sample_standard_normal

end module bem_injection_random