How to add control variable in BLM correctly?

[LiuLianNiuNai]

Hi, thanks for providing this whole package. I have two questions:

First: How to add control variable in BLM correctly?

I'm trying to add some controls in BLM, however, all the example in the documentation is using simulated control variable by tw.categorical_control_params() requiring a1_mu, a1_sig ect.
It's very confusing. Could you provide some clear example about how we add control?

Suppose I want to add the year fixed effect (categorical) firm’s sales (continuous). I have already set up the data using bipartitepandas. Below shows my code.

bdf = bpd.BipartiteDataFrame(df[['i', 'j', 'y', 't', 'time','sales']],
    custom_categorical_dict={'time': True},
    custom_continuous_dict={'sales': True},
    custom_dtype_dict={'time': 'categorical', 'sales': 'continuous'}, 
    custom_how_collapse_dict={'time': 'first','sales': 'first'},
    track_id_changes=True)


clean_params = bpd.clean_params({'connectedness': 'leave_out_spell', 'collapse_at_connectedness_measure': True, 'drop_single_stayers': True})
bdf = bdf.clean(clean_params)


n_firm_types = 6
cluster_params = bpd.cluster_params({'grouping': bpd.grouping.KMeans(n_clusters=n_firm_types)})
bdf = bdf.cluster(cluster_params)
# Convert to event study format
bdf = bdf.to_eventstudy()
# Separate movers and stayers
movers = bdf.get_worker_m()
jdata = bdf.loc[movers, :]
sdata = bdf.loc[~movers, :]
# Initialize BLM variance decomposition estimator
blm_params = tw.blm_params({
      'nk': n_firm_types, 
    'categorical_controls': {'cat_control': 'time'},
    'continuous_controls': {'cts_control': 'sales'}})

blm_fit = tw.BLMVarianceDecomposition(blm_params)

jdata.loc[:, ['w1', 'w2']] = 1
sdata.loc[:, ['w1', 'w2']] = 1

# Fit BLM estimator
blm_fit.fit(jdata, sdata, Q_var=[
                        tw.Q.VarAlpha(),
                        tw.Q.VarPsi(),
                        tw.Q.VarPsiPlusAlpha()
                        tw.Q.VarCovariate(['time', 'sales'])
                    ],
            ncore=48)

Am I right for this?

Second: Can we save the estimated person and frim effect for future use? How to save it?

Best and waiting for your reply