Modified truncate_normal with jax.random.truncated_normal

[mhliu0001]

Previously the truncate_normal distribution was sampled from a normal distribution and then clipping the result. This does not yield a continuous distribution. A more reasonable distribution should be like: https://en.wikipedia.org/wiki/Truncated_normal_distribution

This should affect the microphysics simulation. Maybe the previous version is intentional, so I will let experts comment on which behavior is more correct.

[coveralls]

Pull Request Test Coverage Report for Build 9993829303

Details

• 12 of 14 (85.71%) changed or added relevant lines in 1 file are covered.
• 3 unchanged lines in 1 file lost coverage.
• Overall coverage decreased (-0.1%) to 84.711%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
appletree/randgen.py	12	14	85.71%

Files with Coverage Reduction	New Missed Lines	%
appletree/randgen.py	3	71.03%

Totals
Change from base Build 9897876272:	-0.1%
Covered Lines:	2377
Relevant Lines:	2806

---

💛 - Coveralls

[mhliu0001]

jax-ml/jax#10951
When using jax.random.truncated_normal, we should be careful not to set vmin/vmax too extreme.

But I still think this should be better than using the previous clipping method. The comparison when simulating a truncated normal distribution with mean=0, std=1, vmin=-2, vmax=2 is shown below.

[zihaoxu98]

Thanks! The PR looks good to me

[mhliu0001]

We should revert this PR because NEST also uses the previous truncate_normal in quanta generation. Here NEST defines two types of truncate_normal, one is the non-continuous version (termed rand_gauss with zero_min == true), and the other is the continuous version defined by this PR (termed rand_zero_trunc_gauss). In the quanta generation part, NEST uses rand_gauss so the previous version is more correct.