MNT/ENH: move probabilities tangent mapping code into self-contained …

…functions

- Eliminates duplicate code across problem and results markets.
- Special cases for linear parameters that have zero tangents.

pyblp/markets/market.py

@@ -756,15 +756,69 @@ def compute_utility_derivatives_by_parameter_tangent(
 
         return tangent
 
+    def compute_probabilities_by_parameter_tangent_mapping(
+            self, probabilities: Array, conditionals: Optional[Array], delta: Optional[Array] = None,
+            keep_conditionals: bool = True) -> (
+            Tuple[Dict[int, Array], Dict[int, Optional[Array]]]):
+        """Computing a mapping from parameter index to tangent of probabilities with respect to a parameter. By default,
+        use unchanged delta. By default, also compute conditionals derivatives if computed.
+        """
+        probabilities_mapping: Dict[int, Array] = {}
+        conditionals_mapping: Dict[int, Array] = {}
+        for p, parameter in enumerate(self.parameters.unfixed):
+            probabilities_mapping[p], conditionals_mapping[p] = self.compute_probabilities_by_parameter_tangent(
+                parameter, probabilities, conditionals, delta
+            )
+            if not keep_conditionals:
+                conditionals_mapping[p] = None
+
+        return probabilities_mapping, conditionals_mapping
+
+    def update_probabilities_by_parameter_tangent_mapping(
+            self, probabilities_tangent_mapping: Dict[int, Array],
+            conditionals_tangent_mapping: Dict[int, Optional[Array]], probabilities: Array,
+            conditionals: Optional[Array], xi_jacobian: Array) -> None:
+        """Update tangents of probabilities with respect to parameters to account for the contribution of xi."""
+        probabilities_tensor = conditionals_tensor = None
+        for p, parameter in enumerate(self.parameters.unfixed):
+            probabilities_tangent = probabilities_tangent_mapping[p]
+            conditionals_tangent = conditionals_tangent_mapping[p]
+
+            # total derivatives are zero for beta parameters
+            if isinstance(parameter, BetaParameter):
+                probabilities_tangent[:] = 0
+                if conditionals_tangent is not None:
+                    conditionals_tangent[:] = 0
+                continue
+
+            # derivatives remain zero for gamma parameters
+            if isinstance(parameter, GammaParameter):
+                continue
+
+            # otherwise, need to compute the derivatives of probabilities with respect to xi
+            if probabilities_tensor is None:
+                probabilities_tensor, conditionals_tensor = self.compute_probabilities_by_xi_tensor(
+                    probabilities, conditionals,
+                    compute_conditionals_tensor=any(t is not None for t in conditionals_tangent_mapping.values())
+                )
+
+            # add the contribution of xi
+            probabilities_tangent += np.squeeze(
+                np.moveaxis(probabilities_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
+            )
+            if conditionals_tangent is not None:
+                assert conditionals_tensor is not None
+                conditionals_tangent += np.squeeze(
+                    np.moveaxis(conditionals_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
+                )
+
     def compute_probabilities_by_parameter_tangent(
             self, parameter: Parameter, probabilities: Array, conditionals: Optional[Array],
-            delta: Optional[Array] = None, mu: Optional[Array] = None) -> Tuple[Array, Optional[Array]]:
-        """Compute the tangent of probabilities with respect to a parameter. By default, use unchanged delta and mu."""
+            delta: Optional[Array] = None) -> Tuple[Array, Optional[Array]]:
+        """Compute the tangent of probabilities with respect to a parameter. By default, use unchanged delta."""
         if delta is None:
             assert self.delta is not None
             delta = self.delta
-        if mu is None:
-            mu = self.mu
 
         # without nesting, compute only the tangent of probabilities with respect to the parameter
         if self.H == 0:
@@ -823,7 +877,7 @@ def compute_probabilities_by_parameter_tangent(
             associations = self.groups.expand(group_associations)
 
             # utilities are needed to compute tangents with respect to rho
-            utilities = (delta + mu) / (1 - self.rho)
+            utilities = (delta + self.mu) / (1 - self.rho)
 
             # compute the tangent of conditional probabilities with respect to the parameter
             A = conditionals * utilities / (1 - self.rho)

pyblp/markets/problem_market.py

@@ -46,18 +46,17 @@ def solve(
         if compute_jacobians or moments.MM > 0 or self.K3 > 0:
             probabilities, conditionals = self.compute_probabilities(valid_delta)
 
-        # if needed, pre-compute derivatives of probabilities with respect to parameters
+        # if needed, pre-compute derivatives of probabilities with respect to parameters (conditionals tangents are only
+        #   needed for supply-side Jacobian computation)
         probabilities_tangent_mapping: Dict[int, Array] = {}
         conditionals_tangent_mapping: Dict[int, Array] = {}
         if compute_jacobians:
             assert probabilities is not None
-            for p, parameter in enumerate(self.parameters.unfixed):
-                probabilities_tangent, conditionals_tangent = self.compute_probabilities_by_parameter_tangent(
-                    parameter, probabilities, conditionals, valid_delta
+            probabilities_tangent_mapping, conditionals_tangent_mapping = (
+                self.compute_probabilities_by_parameter_tangent_mapping(
+                    probabilities, conditionals, valid_delta, keep_conditionals=self.K3 > 0
                 )
-                probabilities_tangent_mapping[p] = probabilities_tangent
-                if self.K3 > 0:
-                    conditionals_tangent_mapping[p] = conditionals_tangent
+            )
 
         # compute the Jacobian of xi (equivalently, of delta) with respect to theta
         xi_jacobian = np.full((self.J, self.parameters.P), np.nan, options.dtype)
@@ -71,17 +70,9 @@ def solve(
         # if needed, adjust for the contribution of xi's dependence on theta
         if compute_jacobians and (moments.MM > 0 or self.K3 > 0):
             assert probabilities is not None
-            probabilities_tensor, conditionals_tensor = self.compute_probabilities_by_xi_tensor(
-                probabilities, conditionals, compute_conditionals_tensor=self.K3 > 0
+            self.update_probabilities_by_parameter_tangent_mapping(
+                probabilities_tangent_mapping, conditionals_tangent_mapping, probabilities, conditionals, xi_jacobian
             )
-            for p in range(self.parameters.P):
-                probabilities_tangent_mapping[p] += np.squeeze(
-                    np.moveaxis(probabilities_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
-                )
-                if conditionals_tensor is not None:
-                    conditionals_tangent_mapping[p] += np.squeeze(
-                        np.moveaxis(conditionals_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
-                    )
 
         # compute contributions to micro moments, their Jacobian, and their covariances
         if moments.MM == 0:

pyblp/markets/results_market.py

@@ -1,6 +1,6 @@
 """Market-level structuring of BLP results."""
 
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 import numpy as np
 
@@ -46,17 +46,13 @@ def safely_compute_jacobian_realizations(self, tilde_costs: Array) -> Tuple[Arra
         """
         errors: List[Error] = []
 
-        # if needed, pre-compute probabilities and their derivatives with respect to parameters
+        # pre-compute probabilities and their derivatives with respect to parameters
         probabilities, conditionals = self.compute_probabilities()
-        probabilities_tangent_mapping: Dict[int, Array] = {}
-        conditionals_tangent_mapping: Dict[int, Array] = {}
-        for p, parameter in enumerate(self.parameters.unfixed):
-            probabilities_tangent, conditionals_tangent = self.compute_probabilities_by_parameter_tangent(
-                parameter, probabilities, conditionals
+        probabilities_tangent_mapping, conditionals_tangent_mapping = (
+            self.compute_probabilities_by_parameter_tangent_mapping(
+                probabilities, conditionals, keep_conditionals=self.K3 > 0
             )
-            probabilities_tangent_mapping[p] = probabilities_tangent
-            if self.K3 > 0:
-                conditionals_tangent_mapping[p] = conditionals_tangent
+        )
 
         # compute the demand-side Jacobian
         xi_jacobian, xi_jacobian_errors = self.compute_xi_by_theta_jacobian(
@@ -69,17 +65,9 @@ def safely_compute_jacobian_realizations(self, tilde_costs: Array) -> Tuple[Arra
             omega_jacobian = np.full((self.J, self.parameters.P), np.nan, options.dtype)
         else:
             # adjust for the contribution of xi's dependence on theta
-            probabilities_tensor, conditionals_tensor = self.compute_probabilities_by_xi_tensor(
-                probabilities, conditionals
+            self.update_probabilities_by_parameter_tangent_mapping(
+                probabilities_tangent_mapping, conditionals_tangent_mapping, probabilities, conditionals, xi_jacobian
             )
-            for p in range(self.parameters.P):
-                probabilities_tangent_mapping[p] += np.squeeze(
-                    np.moveaxis(probabilities_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
-                )
-                if conditionals_tensor is not None:
-                    conditionals_tangent_mapping[p] += np.squeeze(
-                        np.moveaxis(conditionals_tensor, 0, 2) @ xi_jacobian[:, [p]], axis=2
-                    )
 
             # compute the supply-side Jacobian
             eta, capital_delta_inverse, eta_errors = self.compute_eta(