Optimize CDF Calculation and Convert NumPy Arrays to Tensors in Benchmark

[yalsaffar]

PR Description

This PR addresses the first step in making AEPsych's functions consistently return PyTorch tensors and expect tensors as input, improving compatibility with GPUs and reducing redundant conversions between NumPy arrays and PyTorch tensors(partially solving #365).

Key changes include:

1. Conversion of np.arrays to tensors in the following files:

   • aepsych/models/base.py:

     • Refactored the p_below_threshold method to operate fully with PyTorch tensors.
     • Replaced norm.cdf() with torch.distributions.Normal(0, 1).cdf() for better GPU compatibility.

   • aepsych/benchmark/problem.py:

     • Significant changes made to ensure consistent use of tensors across the pipeline.
     • The result of f_threshold() now directly returns a PyTorch tensor, ensuring consistency.
     • Additionally, used detach().cpu().numpy() in places where the super().evaluate() method returns float values, ensuring compatibility.

2. Updates in aepsych/tests/test_benchmark.py:

   • Migrated all operations from NumPy to PyTorch.
   • This includes calculations for Brier score and misclassification error, now utilizing torch.mean(), torch.square(), torch.isclose(), and torch.all() to fully align with tensor operations.

Stability:

All test cases have passed successfully in the workflow.

[facebook-github-bot]

@JasonKChow has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[JasonKChow]

There's still some functions that return numpy arrays (despite the documentation specifically saying otherwise). Example: in f_true method of Problem class.

Some typehints are wrong, example f_true method again.

The evaluate method can be rewritten entirely to use torch. Replace pearsonr with corrcoef for torch, you'll need to reformat the data.

[yalsaffar]

Thank you, @JasonKChow, for the valuable feedback!

I’ve made the following updates:

• Updated the evaluate method to fully utilize PyTorch, replacing pearsonr with corrcoef and reformatted the data.
• Fixed the f_true method to ensure it works with tensors and corrected the type hints accordingly.
• The only method still using np.array in Problem.py is sample_y(). Changing this would require significant modifications to the Strategy class, which I plan to address in the next PR!

Please let me know if there is anything I need to change!

[facebook-github-bot]

@JasonKChow has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[JasonKChow]

Why is this cloned and detached? Also why is it converted to tensor if we already assume that self.thresholds will be a tensor?

[yalsaffar]

You're right, there's no need to convert it again, and cloning and detaching are also unnecessary. I'll remove those, test it, and commit the changes.

[facebook-github-bot]

@JasonKChow has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@JasonKChow merged this pull request in 45d8e2d.