diff --git a/src/skillmodels/chs/kalman_filters.py b/src/skillmodels/chs/kalman_filters.py index ad47fd6a..bc316b1e 100644 --- a/src/skillmodels/chs/kalman_filters.py +++ b/src/skillmodels/chs/kalman_filters.py @@ -86,12 +86,19 @@ def kalman_update( ) _residuals = _safe_measurements.reshape(n_obs, 1) - _safe_expected_measurements - _f_stars = jnp.dot(upper_chols, loadings.reshape(n_states, 1)) + # An identity covariance for missing measurements keeps the QR decomposition and + # its derivative well defined; the updated covariance is discarded below. + _safe_upper_chols = jnp.where( + not_missing.reshape(n_obs, 1, 1, 1), + upper_chols, + jnp.eye(n_states, dtype=upper_chols.dtype), + ) + _f_stars = jnp.dot(_safe_upper_chols, loadings.reshape(n_states, 1)) _m = jnp.zeros((n_obs, n_mixtures, n_states + 1, n_states + 1)) _m = _m.at[..., 0, 0].set(meas_sd) _m = _m.at[..., 1:, :1].set(_f_stars) - _m = _m.at[..., 1:, 1:].set(upper_chols) + _m = _m.at[..., 1:, 1:].set(_safe_upper_chols) _r = array_qr_jax(_m)[1] diff --git a/src/skillmodels/chs/likelihood.py b/src/skillmodels/chs/likelihood.py index 463bf7b9..8682574c 100644 --- a/src/skillmodels/chs/likelihood.py +++ b/src/skillmodels/chs/likelihood.py @@ -180,6 +180,7 @@ def log_likelihood_obs( "measurements": measurements, "is_measurement_iteration": is_measurement_iteration, "is_predict_iteration": is_predict_iteration, + "observed_later": get_observation_masks(measurements, iteration_to_period)[0], } _body = functools.partial( @@ -206,6 +207,90 @@ def log_likelihood_obs( ).sum(axis=0) +def get_observation_masks( + measurements: Array | np.ndarray, + iteration_to_period: Array | np.ndarray, +) -> tuple[Array, Array]: + """Return masks of persons with an observed measurement after each update. + + Args: + measurements: Array of shape (n_updates, n_obs). NaN if not observed. + iteration_to_period: Array of length n_updates mapping each update to its + period. Updates of one period are contiguous. + + Returns: + - observed_later: Array of shape (n_updates, n_obs). True if the person has + an observed measurement in a later update. Only these persons are + predicted after the update. + - observed_from_period: Array of shape (n_updates, n_obs). True if the + person has an observed measurement in the period of the update or in a + later period. Filtered states are reported only for these persons. + + """ + observed = jnp.isfinite(measurements) + observed_from = jnp.flip(jnp.cumsum(jnp.flip(observed, axis=0), axis=0) > 0, axis=0) + observed_later = jnp.concatenate( + [observed_from[1:], jnp.zeros_like(observed_from[:1])], axis=0 + ) + periods = jnp.asarray(iteration_to_period) + is_period_start = jnp.concatenate([jnp.array([True]), periods[1:] != periods[:-1]]) + period_start = jax.lax.cummax( + jnp.where(is_period_start, jnp.arange(len(periods)), 0), axis=0 + ) + return observed_later, observed_from[period_start] + + +def predict_if_observed_later( + *, + predict_kwargs: dict[str, Any], + observed_later: Array | np.ndarray, + is_predict_iteration: Array | np.ndarray, + predict_func: Callable, +) -> tuple[Array, Array]: + """Predict the persons with a later observation; keep all others unchanged. + + After the last observed measurement of a person, states and covariances stay at + their last filtered values. The predict step runs on placeholder inputs (zero + states, identity covariance) for these persons and its result is discarded. + The placeholders keep the predict step and its derivatives finite, which a + covariance that is predicted over many unobserved periods does not do. + + Args: + predict_kwargs: Keyword arguments of `predict_func`, with `states` of shape + (n_obs, n_mixtures, n_states) and `upper_chols` of shape + (n_obs, n_mixtures, n_states, n_states). + observed_later: Boolean array of length n_obs. + is_predict_iteration: Whether a predict step follows this update. + predict_func: The predict function. + + Returns: + new_states: Same shape as `predict_kwargs["states"]`. + new_upper_chols: Same shape as `predict_kwargs["upper_chols"]`. + + """ + states = predict_kwargs["states"] + upper_chols = predict_kwargs["upper_chols"] + keep_states = observed_later.reshape(-1, 1, 1) + keep_chols = observed_later.reshape(-1, 1, 1, 1) + safe_kwargs = { + **predict_kwargs, + "states": jnp.where(keep_states, states, 0.0), + "upper_chols": jnp.where( + keep_chols, upper_chols, jnp.eye(states.shape[-1], dtype=upper_chols.dtype) + ), + } + new_states, new_upper_chols = jax.lax.cond( + is_predict_iteration, + lambda kwargs: predict_func(**kwargs), + lambda kwargs: (kwargs["states"], kwargs["upper_chols"]), + safe_kwargs, + ) + return ( + jnp.where(keep_states, new_states, states), + jnp.where(keep_chols, new_upper_chols, upper_chols), + ) + + def _scan_body( carry: dict[str, Array], loop_args: dict[str, Array], @@ -264,16 +349,15 @@ def _scan_body( "observed_factor_vars": observed_factor_vars[t], } - fixed_kwargs = {"predict_func": predict_func} - # ================================================================================== - # Do a predict step or a do-nothing fake predict step + # Do a predict step for persons observed later; freeze all others # ================================================================================== - states, upper_chols, filtered_states = jax.lax.cond( - loop_args["is_predict_iteration"], - functools.partial(_one_arg_predict, **fixed_kwargs), - functools.partial(_one_arg_no_predict, **fixed_kwargs), - predict_kwargs, + filtered_states = states + states, upper_chols = predict_if_observed_later( + predict_kwargs=predict_kwargs, + observed_later=loop_args["observed_later"], + is_predict_iteration=loop_args["is_predict_iteration"], + predict_func=predict_func, ) new_state = { @@ -302,20 +386,3 @@ def _one_arg_anchoring_update( new_log_mixture_weights, new_loglikes, ) - - -def _one_arg_no_predict( - kwargs: dict[str, Any], - predict_func: Callable, # noqa: ARG001 -) -> tuple[Array, Array, Array]: - """Just return the states cond chols without any changes.""" - return kwargs["states"], kwargs["upper_chols"], kwargs["states"] - - -def _one_arg_predict( - kwargs: dict[str, Any], - predict_func: Callable, -) -> tuple[Array, Array, Array]: - """Do a predict step but also return the input states as filtered states.""" - new_states, new_upper_chols = predict_func(**kwargs) - return new_states, new_upper_chols, kwargs["states"] diff --git a/src/skillmodels/chs/likelihood_debug.py b/src/skillmodels/chs/likelihood_debug.py index 88a80481..8a4ccda6 100644 --- a/src/skillmodels/chs/likelihood_debug.py +++ b/src/skillmodels/chs/likelihood_debug.py @@ -11,6 +11,7 @@ from skillmodels.chs.clipping import soft_clipping from skillmodels.chs.kalman_filters_debug import kalman_update +from skillmodels.chs.likelihood import get_observation_masks, predict_if_observed_later from skillmodels.chs.options import CHSEstimationOptions from skillmodels.common.parse_params import parse_params from skillmodels.common.types import ( @@ -90,6 +91,9 @@ def log_likelihood( "log_mixture_weights": log_mixture_weights, } + observed_later, observed_from_period = get_observation_masks( + measurements, iteration_to_period + ) loop_args = { "period": iteration_to_period, "loadings": parsed_params.loadings, @@ -98,6 +102,7 @@ def log_likelihood( "measurements": measurements, "is_measurement_iteration": is_measurement_iteration, "is_predict_iteration": is_predict_iteration, + "observed_later": observed_later, } _body = functools.partial( @@ -146,7 +151,9 @@ def log_likelihood( out["initial_states"] = initial_states out["initial_log_mixture_weights"] = initial_log_mixture_weights - out["filtered_states"] = static_out["states"] + out["filtered_states"] = jnp.where( + observed_from_period[..., None, None], static_out["states"], jnp.nan + ) out["log_mixture_weights"] = static_out["log_mixture_weights"] return out @@ -210,16 +217,15 @@ def _scan_body( "observed_factor_vars": observed_factor_vars[t], } - fixed_kwargs = {"predict_func": predict_func} - # ================================================================================== - # Do a predict step or a do-nothing fake predict step + # Do a predict step for persons observed later; freeze all others # ================================================================================== - states, upper_chols, filtered_states = jax.lax.cond( - loop_args["is_predict_iteration"], - functools.partial(_one_arg_predict, **fixed_kwargs), - functools.partial(_one_arg_no_predict, **fixed_kwargs), - predict_kwargs, + filtered_states = states + states, upper_chols = predict_if_observed_later( + predict_kwargs=predict_kwargs, + observed_later=loop_args["observed_later"], + is_predict_iteration=loop_args["is_predict_iteration"], + predict_func=predict_func, ) new_state = { @@ -249,20 +255,3 @@ def _one_arg_anchoring_update( new_loglikes, debug_info, ) - - -def _one_arg_no_predict( - kwargs: dict[str, Any], - predict_func: Callable[..., tuple[Array, Array]], # noqa: ARG001 -) -> tuple[Array, Array, Array]: - """Just return the states cond chols without any changes.""" - return kwargs["states"], kwargs["upper_chols"], kwargs["states"] - - -def _one_arg_predict( - kwargs: dict[str, Any], - predict_func: Callable[..., tuple[Array, Array]], -) -> tuple[Array, Array, Array]: - """Do a predict step but also return the input states as filtered states.""" - new_states, new_upper_chols = predict_func(**kwargs) - return new_states, new_upper_chols, kwargs["states"] diff --git a/tests/regression_vault/frozen_after_last_observation.json b/tests/regression_vault/frozen_after_last_observation.json new file mode 100644 index 00000000..c734d6de --- /dev/null +++ b/tests/regression_vault/frozen_after_last_observation.json @@ -0,0 +1,636 @@ +{ + "loglikeobs": { + "1": [ + -7.134566985277472, + -55.55718708546832, + -73.88521715757182, + -56.92893771154676, + -63.98411349708587, + -59.35037423918197 + ], + "2": [ + -7.12399477541341, + -57.04178174543407, + -76.30743350276752, + -53.77772210495881, + -62.6912466095914, + -57.28541505388733 + ], + "3": [ + -7.217877801098295, + -54.456874473066875, + -76.4452660397546, + -58.87391428540903, + -66.02071155116758, + -58.5593725153565 + ] + }, + "gradient": { + "1": [ + 4.297177444899113, + -3.814464797323354, + -0.7109773203240399, + 2.34524705915835, + 4.163415239927064, + -1.7326240793409415, + 2.64142239143012, + 1.166524379835459, + 6.2474641806962135, + 0.5496435662161473, + 0.37879282298693884, + 2.5081613053136866, + -7.413440029431087, + 4.662834625196312, + 0.17518320397241105, + 5.7640685121464195, + 1.7166488898871846, + 6.3160201853391165, + -5.29715639260748, + -9.216592398785295, + -4.709923210540689, + 8.685383520656707, + 1.387950023189097, + 13.85723005308087, + -5.932515511113985, + 12.906566695497439, + -4.211739549246215, + -0.3092242645352643, + 1.823369278936776, + -17.3200239629853, + -5.840687526282728, + -6.020940524013762, + -1.9043584360740673, + 0.16477320425738817, + 3.5034313040032785, + -3.757581258172528, + -0.38229637945062933, + 1.9376476867205898, + -0.19075533533785416, + 5.080843061530703, + -1.5815287067571306, + 2.7158085801921965, + 2.672104019934925, + -4.920273977576617, + 2.058181263388754, + -2.8362102359204484, + -1.7613639046302498, + -0.6448799535861616, + 1.018641962817071, + -1.2393551983023168, + 2.889231837692134, + -0.20261239623399607, + -0.1271395076204378, + 0.5504669117028409, + 0.47671248729402443, + -0.9501325523159798, + -0.3552976773114638, + -1.0882706638643591, + -4.116154587884413, + 6.562662691749513, + -1.608954726900272, + 3.5239627998377596, + 6.719967388046421, + 5.35037632037303, + -5.145997278682464, + -1.2434613966816057, + 4.249255250678023, + 4.434174225703568, + 12.649091765443867, + 16.49179860556029, + -7.325033354288369, + 26.048219143699527, + 27.630032372875284, + 17.14034812025027, + 60.48919563208724, + 59.44307576396703, + 20.540960598882634, + 2.8900721469056876, + -0.6447566315768716, + 10.302715548652014, + 18.162714548537593, + 13.755771588735989, + 23.232406952813683, + 18.803641512584814, + 13.709658249555199, + 18.28704121921863, + 4.038871268860952, + 18.88048274424364, + 7.976964114879629, + 26.542125221585003, + 9.89164031152602, + 53.43876042209422, + -4.7293784140143895, + 19.828882318370557, + 44.98346090640962, + 22.430155163279455, + -0.5943689150467011, + -2.2131507886429462, + 0.10201767521436361, + 0.1326195257462195, + 2.3605657209993667, + 0.3825865237975258, + 0.7468290197003896, + 3.482578567297283, + 3.134921287929102, + -0.5817799372026726, + -0.6597317770521076, + 3.2087906369882235, + 11.718203878478526, + 35.73397249015328, + -1.7571103338362546, + -0.38167822547879027, + 0.0, + -5.462024245499286, + 0.39098402104170227, + -3.3815972050311274, + -1.7657798534744291, + -1.6303205976884991, + 0.7473628786612796, + -1.5826864977658957, + 0.4082909008044541, + -2.5805076766536486, + -2.411689178719334, + 1.8245816789758695, + 0.6342589031394174, + -0.9015954448676525, + -0.06431393194574288, + -1.7483993863692644, + -1.9024228257138527, + 2.1685905355528674, + -1.551382432641765, + -0.9519804072113983, + 0.34528956271911354, + -6.181725534264411, + -1.3947529708677777, + 0.6979659483516223, + 1.0235373153278646, + -1.5108650415645148, + 0.4325497350708124, + 0.19255530644254198, + -1.8765149434949107, + 2.9294076348090083, + 0.19993305274001352, + -1.7446419002591371, + 0.5567577470549764, + -3.37294551602771, + -2.1518596338478138, + 1.809437651210114, + 0.26843966115618617, + -0.8329289572865276, + 0.4943889301673793, + -3.7812684898464255, + -1.4715857085215844, + 2.6889633205595542, + 0.38891928200750325, + -3.1155037620136032, + 0.2772906135236465, + -12.246579438681957, + 1.8215522731804932, + 0.9654251771273128, + -1.6848676596093601, + -1.5071967918204388, + -0.624999314986051, + -1.4669839886623863, + 0.8733151707961071, + -3.7668747633244988, + -1.4887734835280975, + -0.48472497603426734, + -0.344347990100532, + -2.4470948420563303, + -0.8287841568931184, + -1.6280570674355133, + -0.8318164687475121, + -0.38823738847660116, + -1.1956784201931314, + -4.26470830091072, + -0.6398944576362511, + -1.5786327370698956, + -1.393337702988892, + 2.097039493739656, + 0.2883432711609144, + -10.351156617159356, + -1.7760870701344669, + -2.4636781265412324, + 0.4255662967473325, + 1.7193135031291085, + 0.5225952240624151, + 5.356163415573095, + -1.3347556450900746, + 0.2620866213347797, + -0.03738939322261937, + -0.2831243778456084, + 0.11390793950114954, + -7.673557357915004, + 2.004354422039077, + 5.314763074491707, + -0.3008072019104723, + -1.9748921915837079, + -1.2563660663291791, + -12.731696664649174 + ], + "2": [ + 6.314972023035169, + -3.3262763397500517, + -0.014333438540929389, + 3.973751875472742, + 4.748227933583724, + -2.531867903342642, + 3.290395685017104, + 1.4971265393532973, + 6.657923275860061, + 0.26104362557898253, + 0.574205922370671, + 1.983852423347992, + -7.112440370983648, + 6.519241063734033, + -0.2615394209582076, + 6.88551981438658, + 1.8176969290519995, + 5.453013819139456, + -6.709977176441288, + -8.462954076501342, + -6.862681557892832, + 8.387622019905727, + 1.1837348792798652, + 8.0283935901493, + -7.720545172683389, + 15.116313099318145, + -5.034389602922248, + -2.5620734285252347, + 0.7942012124978357, + -16.242250233595836, + -3.667778422161512, + -7.731571601284651, + -2.2298246176855083, + -0.19805839407515086, + 3.553118774517982, + -2.7101529800359154, + -0.18190582668382405, + 2.1432717607709435, + -0.2105590908671814, + 2.451201053362255, + -0.009294966197472831, + 1.0678477901329924, + 5.083126631854461, + -4.4020215781041045, + 0.318883581384538, + -3.95514068409182, + -2.189967885428405, + 0.3065904399365461, + 4.059703312392013, + -1.7497498871322295, + 3.1841586463003795, + -1.8723627051663794, + 2.7796204606126724, + -1.3097003551168411, + 1.7859378836900361, + -1.7182077050982592, + 0.5623876785470188, + -0.6081584896854129, + -3.34608891259406, + 4.622636938292654, + -1.303630664752416, + 2.3509552503056934, + 5.925198474218513, + 7.638415839127274, + -4.055548844242049, + -2.471730066409998, + 3.391921898745361, + 6.031314430081508, + 16.473156240989237, + 26.524992532898324, + -7.647761896316289, + 14.444716988529748, + 26.30222708141186, + 22.075992187474316, + 48.38483836795638, + 35.27035529815815, + 19.25417084773588, + 8.353496836510134, + 0.7448366497611791, + 14.466113445433828, + 35.7097394218441, + 16.973781917649646, + 17.190161395767916, + 14.43796405090485, + 27.357677094907977, + 20.140018475640435, + 15.238072048864641, + 6.5471651151354315, + 11.21988377672107, + 38.13769657900284, + 5.775782383615806, + 42.28385782027168, + -4.381070445395332, + 17.93421241921699, + 29.655070141126565, + 30.227229047421147, + -1.272487914079834, + -4.3273153955829775, + 1.8184935432304086, + 0.43281400882512566, + 4.502649259456549, + 1.2819354661011721, + 3.7356337831267266, + 3.912383342221653, + 1.6007179163797611, + -2.589163043355416, + -0.8783649572540031, + 5.6278922811362895, + 13.466758263983387, + 28.146668311095286, + -0.9959875756770216, + -0.6289103788215344, + 0.0, + -5.450309983955263, + 0.38858669497880227, + -3.2305005407599667, + -1.9019842774753206, + -0.6645818412319304, + 0.5578852561876596, + -1.601212185339862, + 0.268070020268566, + -2.1496745490186755, + -3.7760539337125136, + 2.6347619967226663, + 0.261259300145187, + -0.8579141822486528, + 0.25983868877235533, + -1.9095280146617386, + -2.70418280579135, + 2.8256137157528136, + -0.8126081744500454, + -0.6046493568721379, + 0.6438762459352262, + -5.461483839199595, + 0.09757286351535355, + 1.8955518297355216, + 2.235554929948606, + -2.189918965695964, + 1.2060637789328068, + -0.3139892762496771, + -0.17239109189277121, + 3.585847229024873, + -1.262693967876965, + -2.075617177517464, + -0.17076463072793757, + -4.73014077091295, + -2.2879379144425562, + 3.1146637305529175, + 0.02166501530415945, + -0.8030488268014572, + 0.029422780179059818, + -3.2503106381095757, + -1.9642667383357626, + 1.9392908174696333, + 0.13694333701189065, + -2.152280870518295, + 0.3899533351809464, + -13.481147759547206, + 0.6583860820902665, + -1.8953243778776294, + -0.7254484908608814, + -0.2791409381682203, + 0.21431763015414956, + -0.6222123216058854, + 1.4150925812673336, + -3.5679531614558613, + -1.586037143046966, + -0.7841497884600624, + 0.05568973028424721, + -3.5679892287279515, + -0.45559368610363676, + -3.833946994628201, + -0.5737545198900322, + -1.9565725126206976, + -1.3167035094583046, + -5.713419395032774, + -2.773528021317777, + -2.9162452473643166, + -1.6395207188417042, + 2.4470019780374024, + -0.18113226484134212, + -9.222730907232672, + -1.513154754314881, + -1.2893428136078278, + 0.6512430171625737, + 0.7010895433939422, + 0.41636678227409374, + 3.023716035956675, + -0.27936605546929516, + 0.35541728759038915, + -0.913420021892582, + -0.4838287492687128, + 0.12655702175727615, + -8.979472817896385, + 2.1691486202469203, + 5.095816460929513, + -0.6875709977896546, + -1.636225583344514, + -1.2216275213996048, + -10.35366155129594 + ], + "3": [ + 4.143031122979617, + -3.2463983044267746, + -1.5022363453604584, + 3.0695628298170132, + 4.0056069277135755, + -1.9792963851602101, + 2.9765277895004734, + 1.5896734461803903, + 6.629762870201568, + 0.8490761761736145, + 2.268639823065416, + 2.424836307554992, + -11.715890475661809, + 8.887286898588156, + 1.021499005640677, + 6.135207376024319, + 2.266683447692909, + 5.280020447383937, + -10.352475722607664, + -8.893548572797773, + -3.7195003304854057, + 4.850411294813816, + 1.8439232824422647, + 9.882649272265017, + -4.719512769449711, + 16.269028995584637, + -4.971822284871557, + -1.6333609643524802, + -0.40652526831018165, + -10.449765386975736, + -2.928393590735755, + -5.005604553245613, + -2.000676321440141, + -0.33208629369782094, + 3.553239268833788, + -2.510140752389935, + -0.900318978620342, + 2.966135509136266, + 0.23866579445293182, + 2.248809369639784, + -0.3721460516555464, + 0.5444740733172768, + 10.062817596613453, + -6.961903489096512, + 4.995973610220654, + -6.278675495733961, + -2.5201317919255626, + -0.3445946596136129, + 2.5323864356358268, + -2.2742344980447964, + 5.554184175367781, + -2.4275574166325304, + 0.4080299766462159, + 0.034089648908958686, + 1.0584873991410255, + -0.8837732923229139, + -0.019226237510188593, + 0.15725993064445876, + -3.7507237255131214, + 5.239134905715964, + -1.1059233122808143, + 0.18102587743415066, + 3.2216921808364214, + 5.395401956391401, + -4.631018098165277, + -0.9040551765487516, + 1.034364846159453, + 9.685665007591595, + 10.229161155469164, + 21.575538866394787, + -7.742117047408004, + 16.227608517902073, + 25.931917232094843, + 17.152963412571935, + 97.69103239347658, + 50.750371888288385, + 47.17519386519554, + 18.192504601440895, + 3.1095113521680564, + 11.247146677689603, + 25.85990461835221, + 19.326704490357585, + 29.732556270788507, + 15.144907889988001, + 9.00911587435589, + 8.190097643897472, + 7.26172872207615, + 9.407740238978908, + 4.0765161328756365, + 38.79098595561339, + 8.968674915782543, + 46.912007327877106, + -4.34713826258746, + 7.886704043516098, + 26.30488932465959, + 18.08489794028683, + -0.39209225993963615, + -2.073374799246762, + 0.5784952472641549, + 3.3994152817323497, + 4.579175903026836, + 0.8463947557867268, + 2.244835917990992, + 5.428231059087421, + 2.199538568387735, + -3.0495204406096046, + -0.7271247439343255, + 3.46842081039895, + 8.716378105757633, + 21.930959070221906, + -1.5172928754092327, + -0.4699405641088104, + 0.0, + -5.356632556415679, + 0.5251340744150327, + -3.485782317154707, + -1.6379880548996053, + -0.5870884560598393, + 0.5383944299837999, + -1.8274576917957135, + -0.3062960687399489, + -1.7808143811956523, + -3.2212664750640205, + 1.2854170037422044, + 0.9391852147999148, + -0.7504612816867723, + -0.17685548117726219, + -0.8766452877339241, + -2.1075388632853462, + 3.435384258818232, + -0.8649623127266587, + -1.276949399201559, + 0.18764059609566108, + -6.868121286545322, + -1.1348100206417058, + 1.4856773930218174, + 1.7952583224035827, + -1.0952823646216894, + 1.2712402931261932, + -0.9025410994690268, + -0.43659763599525026, + 4.5898232075320315, + -0.8880203455898975, + -2.459285477828683, + -0.10931719890125723, + -5.503595950861315, + -1.7444949765163689, + 2.816089798083373, + 0.005218000318367466, + -1.2504667753108285, + 0.1232359655341599, + -3.7435599472578924, + -1.6075754641709714, + 2.422427538308148, + 0.5700118414254053, + -2.349670893580335, + -0.10803750217703381, + -11.836155423341232, + 1.521691356328898, + -2.149200548794358, + -1.2656700309599345, + 0.8129034245046496, + 0.22639398773666333, + 0.43843226375597627, + 1.9056978255008559, + -3.440275198013756, + -1.6964939691848295, + -0.18999519870043186, + 0.10420168700358229, + -1.9364178343431506, + 0.30990661472797976, + -3.9698277571976965, + -0.45447922452143436, + -1.9348026983039988, + -1.1003146875807508, + -6.251264116378284, + -2.300614927912126, + -2.036961139036348, + -0.8279722701040468, + 1.2117845111955245, + 0.2295623948876904, + -11.236023932568232, + -1.4852216401167715, + -2.147653751966572, + 0.7179747556106842, + 1.001118197755333, + 0.4243399191261273, + 4.079881095288144, + -0.3891413833113221, + -0.7455636001432385, + -0.5132679024790967, + -0.1413140127255705, + -0.07876034372789575, + -7.454068563459153, + 1.5915058024602475, + 4.286620668689974, + -1.0087940120031775, + -2.027534451143201, + -0.6941838840441863, + -8.363547822027073 + ] + } +} diff --git a/tests/test_frozen_after_last_observation.py b/tests/test_frozen_after_last_observation.py new file mode 100644 index 00000000..5b44e44f --- /dev/null +++ b/tests/test_frozen_after_last_observation.py @@ -0,0 +1,147 @@ +"""The CHS filter stops predicting a person after their last observed measurement. + +Fixture: a two-factor translog panel over eight periods. Person 0 is last observed +in period 1 (death or attrition); person 1 has a gap in period 3 and is observed +again afterwards. With large squared transition terms, the predicted covariance of +person 0 explodes in the periods after the last observation. +""" + +import json +from pathlib import Path + +import jax +import numpy as np +import pandas as pd +import pytest +from numpy.testing import assert_allclose + +from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions +from skillmodels.common.model_spec import FactorSpec, ModelSpec + +jax.config.update("jax_enable_x64", True) + +LATENT = ("fac1", "fac2") +N_PERIODS = 8 +N_PERSONS = 6 +LAST_OBSERVED_PERIOD = {0: 1} +GAP = (1, 3) +REFERENCE_FILE = ( + Path(__file__).parent / "regression_vault" / "frozen_after_last_observation.json" +) + + +def _model() -> ModelSpec: + factors = { + f: FactorSpec( + measurements=((f"{f}_m1", f"{f}_m2"),) * N_PERIODS, + transition_function="translog", + ) + for f in LATENT + } + return ModelSpec(factors=factors) + + +def _data() -> pd.DataFrame: + rng = np.random.default_rng(0) + idx = pd.MultiIndex.from_product( + [range(N_PERSONS), range(N_PERIODS)], names=["id", "period"] + ) + cols = [f"{f}_m{j}" for f in LATENT for j in (1, 2)] + df = pd.DataFrame(rng.normal(size=(len(idx), len(cols))), index=idx, columns=cols) + for person, period in LAST_OBSERVED_PERIOD.items(): + df.loc[(person, slice(period + 1, None)), :] = np.nan + df.loc[GAP, :] = np.nan + return df + + +@pytest.fixture(scope="module") +def inputs() -> dict: + return get_maximization_inputs( + _model(), + _data(), + chs_options=CHSEstimationOptions(start_params_strategy="none"), + ) + + +def _params(template: pd.DataFrame, squares: float, seed: int | None) -> pd.DataFrame: + """Fill the template; squared transition terms get the value `squares`. + + With a `seed`, every free value is perturbed by a uniform factor in [0.8, 1.2]. + """ + values = [] + for category, _, name1, name2 in template.index: + if category == "initial_cholcovs": + row, col = name2.split("-") + val = 1.0 if row == col else 0.2 + elif category == "transition": + if "**" in name2: + val = squares + elif "*" in name2: + val = 0.1 + else: + val = 1.1 if name2 == name1 else 0.0 + else: + val = { + "controls": 0.1, + "loadings": 1.0, + "meas_sds": 0.5, + "shock_sds": 0.3, + "initial_states": 0.0, + "mixture_weights": 1.0, + }[category] + values.append(val) + params = template.copy() + params["value"] = values + if seed is not None: + rng = np.random.default_rng(seed) + free = template["value"].isna().to_numpy() + params.loc[free, "value"] *= rng.uniform(0.8, 1.2, size=free.sum()) + return params + + +def test_gradient_matches_finite_differences_with_exploding_post_exit_prediction( + inputs: dict, +) -> None: + params = _params(inputs["params_template"], squares=1.0, seed=None) + _, grad = inputs["loglike_and_gradient"](params) + free = inputs["params_template"]["value"].isna().to_numpy() + step = 1e-6 + numerical = [] + for i in np.flatnonzero(free): + up, down = params.copy(), params.copy() + up.iloc[i, up.columns.get_loc("value")] += step + down.iloc[i, down.columns.get_loc("value")] -= step + numerical.append( + (float(inputs["loglike"](up)) - float(inputs["loglike"](down))) / (2 * step) + ) + assert_allclose(np.asarray(grad)[free], numerical, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize("seed", [1, 2, 3]) +def test_loglikeobs_equal_reference_values(inputs: dict, seed: int) -> None: + params = _params(inputs["params_template"], squares=0.6, seed=seed) + expected = json.loads(REFERENCE_FILE.read_text())["loglikeobs"][str(seed)] + assert_allclose(np.asarray(inputs["loglikeobs"](params)), expected, rtol=1e-12) + + +@pytest.mark.parametrize("seed", [1, 2, 3]) +def test_gradient_equals_reference_values_where_finite(inputs: dict, seed: int) -> None: + params = _params(inputs["params_template"], squares=0.6, seed=seed) + expected = json.loads(REFERENCE_FILE.read_text())["gradient"][str(seed)] + _, grad = inputs["loglike_and_gradient"](params) + assert_allclose(np.asarray(grad), expected, rtol=1e-10, atol=1e-12) + + +def test_filtered_states_are_nan_exactly_after_last_observation(inputs: dict) -> None: + params = _params(inputs["params_template"], squares=0.6, seed=None) + states = inputs["debug_loglike"](params)["filtered_states"] + is_nan = states.set_index(["id", "aug_period"])[list(LATENT)].isna().all(axis=1) + expected = pd.Series( + [ + person in LAST_OBSERVED_PERIOD and period > LAST_OBSERVED_PERIOD[person] + for person, period in is_nan.index + ], + index=is_nan.index, + ) + pd.testing.assert_series_equal(is_nan, expected, check_names=False)