diffusion_coefficient function change

[YitongPan1]

Summary

Major changes:

• feature 1: ...
• fix 1: ...

Todos

If this is work in progress, what else needs to be done?

• feature 2: ...
• fix 2:

Checklist

• Google format doc strings added.
• Code linted with ruff. (For guidance in fixing rule violates, see rule list)
• Type annotations included. Check with mypy.
• Tests added for new features/fixes.
• I have run the tests locally and they passed.

Tip: Install pre-commit hooks to auto-check types and linting before every commit:

pip install -U pre-commit
pre-commit install

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.79%. Comparing base (98e7824) to head (3f0277a).
Report is 38 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #215   +/-   ##
=======================================
  Coverage   81.79%   81.79%           
=======================================
  Files           9        9           
  Lines        1494     1494           
  Branches      256      257    +1     
=======================================
  Hits         1222     1222           
  Misses        226      226           
  Partials       46       46           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.