Add Type Hints, Fix Mypy Errors, and Refactor Code in AEPsych Directory #413

yalsaffar · 2024-10-22T16:48:24Z

Description:

This PR adds type hints, resolves Mypy errors, and refactors code across multiple modules in AEPsych. Key changes include type hint additions, tensor conversions for compatibility, and improved error handling. This is the first step towards resolving issue #366.

Changes by Directory:

1. `acquisitions/`:

Added type hints across files.
Refactored construct_inputs_global_lookahead to use tensors for lb and ub instead of lists.

2. `benchmark/`:

Added type hints.
In problem.py, updated LSEProblemWithEdgeLogging.evaluate to use tensors for metrics["prop_edge_sampling_err"].

3. `database/`, `generators/`, `kernels/`, `likelihoods/`:

Added type hints throughout.
kernels/pairwisekernel.py: Updated PairwiseKernel.forward to avoid assigning tensor d as both an integer and tensor, ensuring Mypy compatibility.

4. `models/`:

Added type hints.
gp_classification.py: Added assertions (inducing_points is not None) in the constructor and _reset_variational_strategy.
semi_p.py: Added tensor conversion for X in SemiParametricGPModel.posterior to avoid Mypy errors.

5. Root-Level AEPsych Files:

Added type hints throughout.
strategy.py: Added assertions to check for None in self.model, self.x, and self.y to prevent errors.
utils.py: Wrapped all occurrences of parname with str() for compatibility with the Config class.
plotting.py:
- Refactored plot_strat_3d to improve handling of contour_levels:
  - Introduced detailed checks for contour_levels_list to ensure it is Sized (e.g., list or array).
  - Improved type handling for slice_vals, ensuring proper error messages and type checking.
  - Added error handling for cases where the model is None, raising a RuntimeError instead of silently failing.
- Updated plot_slice to ensure better control over contour_levels:
  - Enforced type checking on contour_levels and ensured non-None values are provided, with clearer error messages.
  - Refactored the function to raise an explicit error (ValueError) if contour_levels is incorrectly set to None.

JasonKChow

Some typehints are missing:
from_config all over are missing return hints
lookahead_type
init missing return None -- Semi_p.py file
problem.py Problem class properties missing return typehints.
problem.py has a TODO comment in sample_y method, maybe make a small PR to fix this if possible.
config.py Config missing type hints in methods (e.g., get_section())
db.py Database missing type hints
Likelihood semi_p.py file "expected_log_prob" missing type hints.

Shouldn't everything be using Tensors now?

aepsych/acquisition/lookahead.py

aepsych/acquisition/lse.py

aepsych/acquisition/mc_posterior_variance.py

aepsych/acquisition/mutual_information.py

aepsych/benchmark/problem.py

aepsych/utils.py

yalsaffar · 2024-10-24T20:42:50Z

I’ve gone ahead and added return type hints for the from_config method throughout the codebase. Along the way, I also fixed the type hints for lookahead_type, and made sure all the __init__ methods that were missing None return types are now sorted out.

The TODO comment in Problem.py is still on the radar, but I’ll tackle that in a follow-up PR after this one. I also worked on config.py by fixing the missing type hints in methods like get_section(), and I cleaned up some issues in db.py related to the Database.

For semi_p.py, I added the missing type hints for expected_log_prob, even though it’s not currently being used. Most of the specific changes have already been resolved, but a few are still up for discussion. As soon as I get more feedback, I’ll add those as well.

Lastly, most parts of the code where tensors are needed for GPU support have been converted. The one exception is in the add_data method, which still accepts both tensors and np.ndarrays. This is because, as discussed in #403, it's the first point where data is accepted, and from there, everything is processed as tensors.

JasonKChow

Many comments.

In general: avoid adding extra asserts to fix mypy problems, instead work out how to make the type hints align. Basically, no actual functional code should be written for most mypy problems, just modifying hints.

aepsych/utils.py

aepsych/database/tables.py

JasonKChow · 2024-10-25T00:38:49Z

aepsych/generators/acqf_thompson_sampler_generator.py

@@ -61,7 +61,7 @@ def __init__(
        self.samps = samps
        self.stimuli_per_trial = stimuli_per_trial

-    def _instantiate_acquisition_fn(self, model: ModelProtocol):
+    def _instantiate_acquisition_fn(self, model: ModelProtocol) -> AnalyticExpectedUtilityOfBestOption:


Double check again. Doesn't appear changed.

aepsych/generators/monotonic_rejection_generator.py

JasonKChow · 2024-10-25T00:43:03Z

aepsych/models/gp_classification.py

        inducing_points = select_inducing_points(
            inducing_size=self.inducing_size,
            covar_module=self.covar_module,
            X=self.train_inputs[0],
            bounds=self.bounds,
            method=self.inducing_point_method,
        )
+        assert inducing_points is not None, "Inducing points cannot be None"


These are presumably for mypy errors, I don't like adding asserts to solve mypy errors, see if you can figure out a different way.

JasonKChow · 2024-10-25T00:45:25Z

aepsych/models/utils.py

@@ -57,7 +58,7 @@ def select_inducing_points(
    X: Optional[torch.Tensor] = None,
    bounds: Optional[Union[torch.Tensor, np.ndarray]] = None,
    method: str = "auto",
-):
+) -> Optional[torch.Tensor]:


Should always be tensors (also this Optional might be the cause of other problems above that require Asserts).

JasonKChow · 2024-10-25T00:47:32Z

aepsych/plotting.py

+    contour_levels_list: List[float] = []
+
+    if contour_levels is True:
+        contour_levels_list = [0.75]


It's a very complex argument it seems. It has to allow bool to default it to all ints. Double check if this is absolutely necessary.

JasonKChow · 2024-10-25T00:49:30Z

aepsych/strategy.py

        self.is_finished = True

    @property
-    def finished(self):
+    def finished(self) -> Union[bool,torch.Tensor, None]:


This should be strictly bool, why not?

JasonKChow · 2024-10-25T00:51:32Z

aepsych/strategy.py

        if self.can_fit:
            if self.keep_most_recent is not None:
                try:
+                    assert self.model is not None, "a model is needed here!"


Assert probably due to above Optional on can_fit. Wherever assert ends up being weirdly needed, it's likely because of an odd typing somewhere else that should be fixed instead.

JasonKChow · 2024-10-25T00:57:05Z

aepsych/models/ordinal_gp.py

Add type hints for arguments, then try again, if not ask again.

facebook-github-bot · 2024-10-25T23:03:00Z

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

aepsych/acquisition/lookahead.py

aepsych/benchmark/problem.py

aepsych/kernels/pairwisekernel.py

aepsych/models/semi_p.py

aepsych/strategy.py

facebook-github-bot · 2024-10-29T16:37:51Z

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

facebook-github-bot · 2024-10-29T22:56:11Z

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

facebook-github-bot · 2024-10-29T23:39:42Z

@JasonKChow merged this pull request in fba7f99.

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 22, 2024

JasonKChow requested changes Oct 23, 2024

View reviewed changes

JasonKChow requested changes Oct 25, 2024

View reviewed changes

JasonKChow approved these changes Oct 25, 2024

View reviewed changes

crasanders suggested changes Oct 28, 2024

View reviewed changes

yalsaffar added 15 commits October 30, 2024 01:50

add type hints to acquisition files

a383ea4

adding type hints to benchmark files

3d57415

adding type hints to database files

48883ae

adding type hints to generators files

902ad7f

adding type hints to kernels files

9e686ae

fix an import in pairwisekernal

df1b0c0

adding type hints to likelihoods files

daf3e2e

adding type hints to models files

e212c8c

partially adding type hints to root aepsych files

f9eb7b5

fixing mypy errors

e4852b9

remove branch from workflow

181153b

fix comments in plotting

4d3a16e

fix the changes requested

e152b51

fixing the needed changes

69aa9f0

fix proposed changes, mainly insure robust error handling and type hints

9eb3cd3

yalsaffar force-pushed the fix-types branch from d63ae86 to 9eb3cd3 Compare October 29, 2024 22:52

facebook-github-bot closed this in fba7f99 Oct 29, 2024

facebook-github-bot added the Merged label Oct 29, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Type Hints, Fix Mypy Errors, and Refactor Code in AEPsych Directory #413

Add Type Hints, Fix Mypy Errors, and Refactor Code in AEPsych Directory #413

yalsaffar commented Oct 22, 2024

JasonKChow left a comment

yalsaffar commented Oct 24, 2024

JasonKChow left a comment

JasonKChow Oct 25, 2024

JasonKChow Oct 25, 2024

JasonKChow Oct 25, 2024

JasonKChow Oct 25, 2024

JasonKChow Oct 25, 2024

JasonKChow Oct 25, 2024

JasonKChow Oct 25, 2024

facebook-github-bot commented Oct 25, 2024

facebook-github-bot commented Oct 29, 2024

facebook-github-bot commented Oct 29, 2024

facebook-github-bot commented Oct 29, 2024

Add Type Hints, Fix Mypy Errors, and Refactor Code in AEPsych Directory #413

Add Type Hints, Fix Mypy Errors, and Refactor Code in AEPsych Directory #413

Conversation

yalsaffar commented Oct 22, 2024

Description:

Changes by Directory:

1. acquisitions/:

2. benchmark/:

3. database/, generators/, kernels/, likelihoods/:

4. models/:

5. Root-Level AEPsych Files:

JasonKChow left a comment

Choose a reason for hiding this comment

yalsaffar commented Oct 24, 2024

JasonKChow left a comment

Choose a reason for hiding this comment

JasonKChow Oct 25, 2024

Choose a reason for hiding this comment

JasonKChow Oct 25, 2024

Choose a reason for hiding this comment

JasonKChow Oct 25, 2024

Choose a reason for hiding this comment

JasonKChow Oct 25, 2024

Choose a reason for hiding this comment

JasonKChow Oct 25, 2024

Choose a reason for hiding this comment

JasonKChow Oct 25, 2024

Choose a reason for hiding this comment

JasonKChow Oct 25, 2024

Choose a reason for hiding this comment

facebook-github-bot commented Oct 25, 2024

facebook-github-bot commented Oct 29, 2024

facebook-github-bot commented Oct 29, 2024

facebook-github-bot commented Oct 29, 2024

1. `acquisitions/`:

2. `benchmark/`:

3. `database/`, `generators/`, `kernels/`, `likelihoods/`:

4. `models/`: