InferenceGym: Add log-return S&P500 VectorizedStochasticVolatilityModel.

PiperOrigin-RevId: 379575684

Author: SiegeLordEx

spinoffs/inference_gym/inference_gym/internal/data.py

@@ -30,7 +30,8 @@
     'convection_lorenz_bridge',
     'german_credit_numeric',
     'radon',
-    'sp500_closing_prices',
+    'sp500_returns',
+    'sp500_log_returns',
     'synthetic_item_response_theory',
     'synthetic_log_gaussian_cox_process',
     'synthetic_plasma_spectroscopy',
@@ -427,9 +428,12 @@ def load_dataset():
   )
 
 
-def sp500_closing_prices(num_points=None):
+def sp500_returns(num_points=None):
   """Dataset of mean-adjusted returns of the S&P 500 index.
 
+  The returns here are the absolute returns (current price minus previous
+  price).
+
   Each of the 2516 entries represents the adjusted return of the daily closing
   price relative to the previous close, for each (non-holiday) weekday,
   beginning 6/26/2010 and ending 6/24/2020.
@@ -448,6 +452,31 @@ def sp500_closing_prices(num_points=None):
       centered_returns=returns[-num_points:] - np.mean(returns[-num_points:]))
 
 
+def sp500_log_returns(num_points=None):
+  """Dataset of mean-adjusted log returns of the S&P 500 index.
+
+  A log return is the change in the log-transformed closing prices, which also
+  corresponds to the log of relative returns (current price divided by previous
+  price).
+
+  Each of the 2516 entries represents the adjusted log return of the daily
+  closing price relative to the previous close, for each (non-holiday) weekday,
+  beginning 6/26/2010 and ending 6/24/2020.
+
+  Args:
+    num_points: Optional `int` length of the series to return. If specified,
+      only the final `num_points` returns are centered and returned.
+      Default value: `None`.
+  Returns:
+    dataset: A Dict with the following keys:
+      `centered_returns`: float `Tensor` daily returns, minus the mean return.
+  """
+  returns = np.diff(np.log(sp500_closing_prices_lib.CLOSING_PRICES))
+  num_points = num_points or len(returns)
+  return dict(
+      centered_returns=returns[-num_points:] - np.mean(returns[-num_points:]))
+
+
 def synthetic_item_response_theory(
     train_fraction=1.,
     shuffle=True,

spinoffs/inference_gym/inference_gym/internal/data_test.py

@@ -60,7 +60,15 @@ def load_dataset_in_xla():
   def testStochasticVolatilityModelSP500(self):
     num_train_points = 2516
 
-    dataset = data.sp500_closing_prices()
+    dataset = data.sp500_returns()
+
+    self.assertEqual((num_train_points,), dataset['centered_returns'].shape)
+    self.assertAllClose(0.0, np.mean(dataset['centered_returns']), atol=1e-5)
+
+  def testStochasticVolatilityModelLogSP500(self):
+    num_train_points = 2516
+
+    dataset = data.sp500_log_returns()
 
     self.assertEqual((num_train_points,), dataset['centered_returns'].shape)
     self.assertAllClose(0.0, np.mean(dataset['centered_returns']), atol=1e-5)

spinoffs/inference_gym/inference_gym/targets/BUILD

@@ -551,6 +551,8 @@ py_library(
         ":bayesian_model",
         ":model",
         "//inference_gym/internal:data",
+        "//inference_gym/targets/ground_truth:stochastic_volatility_log_sp500",
+        "//inference_gym/targets/ground_truth:stochastic_volatility_log_sp500_small",
         "//inference_gym/targets/ground_truth:stochastic_volatility_sp500",
         "//inference_gym/targets/ground_truth:stochastic_volatility_sp500_small",
     ],

spinoffs/inference_gym/inference_gym/targets/__init__.py

@@ -51,6 +51,8 @@
   from inference_gym.targets.stochastic_volatility import StochasticVolatilitySP500
   from inference_gym.targets.stochastic_volatility import StochasticVolatilitySP500Small
   from inference_gym.targets.vectorized_stochastic_volatility import VectorizedStochasticVolatility
+  from inference_gym.targets.vectorized_stochastic_volatility import VectorizedStochasticVolatilityLogSP500
+  from inference_gym.targets.vectorized_stochastic_volatility import VectorizedStochasticVolatilityLogSP500Small
   from inference_gym.targets.vectorized_stochastic_volatility import VectorizedStochasticVolatilitySP500
   from inference_gym.targets.vectorized_stochastic_volatility import VectorizedStochasticVolatilitySP500Small
   from inference_gym.targets.vector_model import VectorModel
@@ -90,6 +92,8 @@
     'SyntheticLogGaussianCoxProcess',
     'SyntheticPlasmaSpectroscopy',
     'VectorizedStochasticVolatility',
+    'VectorizedStochasticVolatilityLogSP500',
+    'VectorizedStochasticVolatilityLogSP500Small',
     'VectorizedStochasticVolatilitySP500',
     'VectorizedStochasticVolatilitySP500Small',
     'VectorModel',

spinoffs/inference_gym/inference_gym/targets/ground_truth/BUILD

@@ -126,6 +126,26 @@ py_library(
     ],
 )
 
+# pytype_strict
+py_library(
+    name = "stochastic_volatility_log_sp500",
+    srcs = ["stochastic_volatility_log_sp500.py"],
+    srcs_version = "PY3",
+    deps = [
+        # numpy dep,
+    ],
+)
+
+# pytype_strict
+py_library(
+    name = "stochastic_volatility_log_sp500_small",
+    srcs = ["stochastic_volatility_log_sp500_small.py"],
+    srcs_version = "PY3",
+    deps = [
+        # numpy dep,
+    ],
+)
+
 # pytype_strict
 py_library(
     name = "stochastic_volatility_sp500",