Add lognormal_dist_from_mean_std helper function.

PiperOrigin-RevId: 803606606

Author: ahalev

CHANGELOG.md

@@ -23,6 +23,8 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+* Add `lognormal_dist_from_mean_std` helper function.
+
 ## [1.2.0] - 2025-09-04
 
 * Fix channel data misalignment in `Analyzer.hill_curves` when input channels

meridian/model/prior_distribution.py

@@ -34,6 +34,7 @@
     'IndependentMultivariateDistribution',
     'PriorDistribution',
     'distributions_are_equal',
+    'lognormal_dist_from_mean_std',
 ]
 
 
@@ -1173,6 +1174,32 @@ def distributions_are_equal(
   return True
 
 
+def lognormal_dist_from_mean_std(
+    mean: float | Sequence[float], std: float | Sequence[float]
+) -> backend.tfd.LogNormal:
+  """Define a lognormal distribution from its mean and standard deviation.
+
+  This function parameterizes lognormal distributions by their mean and
+  standard deviation.
+
+  Args:
+    mean: A positive float or array-like object defining the distribution mean.
+    std: A non-negative float or array-like object defining the distribution
+      standard deviation.
+
+  Returns:
+    A `backend.tfd.LogNormal` object with the input mean and standard deviation.
+  """
+
+  mean = np.asarray(mean)
+  std = np.asarray(std)
+
+  mu = np.log(mean) - 0.5 * np.log((std / mean) ** 2 + 1)
+  sigma = np.sqrt(np.log((std / mean) ** 2 + 1))
+
+  return backend.tfd.LogNormal(mu, sigma)
+
+
 def _convert_to_deterministic_0_distribution(
     distribution: backend.tfd.Distribution,
 ) -> backend.tfd.Distribution:

meridian/model/prior_distribution_test.py

@@ -1898,5 +1898,44 @@ def test_independent_distributions_variance(
     np.testing.assert_allclose(variance, expected_variance, rtol=1e-5)
 
 
+class TestLognormalDistFromMeanStd(parameterized.TestCase):
+
+  @parameterized.product(
+      mean=(1.0, 2.0, 3.0),
+      std=(1.0, 2.0, 3.0),
+      input_type=(float, int, np.float32, backend.to_tensor)
+  )
+  def test_correct_mean_std_scalar(
+      self, mean, std, input_type
+  ):
+    mean = input_type(mean)
+    std = input_type(std)
+
+    dist = prior_distribution.lognormal_dist_from_mean_std(mean, std)
+
+    np.testing.assert_allclose(dist.mean(), mean, rtol=1e-5)
+    np.testing.assert_allclose(dist.stddev(), std, rtol=1e-5)
+
+  @parameterized.product(
+      mean=((1.0,), (1.0, 2.0,)),
+      std=((2.0,), (2.0, 3.0,)),
+      input_type=(tuple, list, np.array, backend.to_tensor)
+  )
+  def test_correct_mean_std_array(
+      self, mean, std, input_type
+  ):
+    mean = input_type(mean)
+    std = input_type(std)
+
+    dist = prior_distribution.lognormal_dist_from_mean_std(mean, std)
+
+    expected_len = max(len(mean), len(std))
+    expected_mean = np.broadcast_to(mean, expected_len)
+    expected_std = np.broadcast_to(std, expected_len)
+
+    np.testing.assert_allclose(dist.mean(), expected_mean, rtol=1e-5)
+    np.testing.assert_allclose(dist.stddev(), expected_std, rtol=1e-5)
+
+
 if __name__ == '__main__':
   absltest.main()