<jupyter_start><jupyter_text>Performance penalty by using fewer taps<jupyter_code>%pip install --upgrade https://github.com/remifan/commplax/archive/master.zip
%pip install --upgrade https://github.com/remifan/LabPtPTm2/archive/master.zip
%pip install --upgrade https://github.com/remifan/gdbp_study/archive/master.zip
import os
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
from functools import partial
import matplotlib.pyplot as plt
from commplax import util
from gdbp import gdbp_base as gb, data as gdat, plot as gplt, aux
ds_train, ds_test = gdat.load(1, 0, 4, 2)[0], gdat.load(2, 0, 4, 1)[0]
aux.dl_trained_params() # download pre-trained parameters
loc_trained_params = './trained_params/pretrained'
# use Pandas.Dataframe to store the results
df_test_res = pd.DataFrame({c: pd.Series(dtype=t) for c, t in {'ntaps': 'int',
                                                               'dtaps': 'int',
                                                               'xi': 'float',
                                                               'Model': 'str',
                                                               'Q': 'float'}.items()})
def init_models(data, hparams, mode='train'):
  model_conf = {'mode':  mode,
                'dtaps': hparams['dtaps'],
                'ntaps': hparams['ntaps'],
                'init_fn': gb.fdbp_init(data.a, steps=3, xi=hparams['xi'])}
  model_init = partial(gb.model_init, data, model_conf)
  fdbp = model_init([('fdbp_0',)]            , name='FDBP')
  edbp = model_init([('fdbp_0', r'DConv_\d')], name='EDBP')
  gdbp = model_init([],                        name='GDBP')
  return fdbp, edbp, gdbp


def sweep_hparams(hparams, df_test_res=df_test_res, use_pretrained_params=False, save_params=False):
  for i, hp in enumerate(tqdm(hparams, desc='sweeping hparams')):  # iterate hparams
    models_train = init_models(ds_train, hp)
    models_test = init_models(ds_test, hp, mode='test')
    
    for j, (model_train, model_test) in enumerate(tqdm(zip(models_train, models_test),
                                                       total=len(models_train),
                                                       desc='iterating models',
                                                       leave=False)):  # iterate models
      params_file = os.path.join(loc_trained_params,
                                 'snr_vs_dtaps_ntaps',
                                 'params_%d_%d' % (i, j)) # params_{hparam}_{model}
      
      if use_pretrained_params:
        params = util.load_variable(params_file)
      else:
        params_queue = [None] * 3
        
        for _, p, _ in gb.train(model_train, ds_train, n_iter=2000):
          params_queue.append(p)
          params = params_queue.pop(0)
        
        if save_params:
          util.save_variable(params, params_file)
        
      test_Q = gb.test(model_test, params, ds_test)[0].QSq.total
      # collect result
      df_test_res = df_test_res.append({'dtaps': hp['dtaps'],
                                        'ntaps': hp['ntaps'],
                                        'xi': hp['xi'],
                                        'Model': model_test.name,
                                        'Q': test_Q},
                                       ignore_index=True)
    
    if i % 10 == 0:
      util.clear_xla_cache() # clear compile cache periodically to save RAM

  return df_test_res
# make sweeping grid

# caveats: taps must be odd number
NTAPS = np.arange(1, 41, 6)
DTAPS = np.arange(181, 261, 12)
OPTXI = pd.read_csv('../../source_data/fdbp_opt_xi_ntaps.csv').values[:, 0]  # read optimal xi given dtaps

hparams = []
for ntaps, xi in zip(NTAPS, OPTXI):
  for dtaps in DTAPS:
    hparams.append({'ntaps': ntaps, 'dtaps': dtaps, 'xi': xi})
# it may take a while to finish
df_test_res = sweep_hparams(hparams, df_test_res, use_pretrained_params=True, save_params=False)
df_test_res
fig, ax = plt.subplots(subplot_kw={"projection": "3d"}, figsize=(5, 5), dpi=200)
ax.view_init(19, -30)
curves = 20
cms = [plt.cm.Greys, plt.cm.Blues, plt.cm.Reds]

colors = plt.cm.RdPu_r(np.linspace(0., 0.5, 3))[::-1]

for cm, mod in zip(cms, ['FDBP', 'EDBP', 'GDBP']):
  df = df_test_res.groupby('Model').get_group(mod).pivot('dtaps', 'ntaps', 'Q')
  (X, Y), Z = np.meshgrid(df.columns.values, df.index.values), df.values
  gplt.wireframe_cmap(ax, X, Y, Z, cmap=cm, offset=[-1.2, 0], label=mod)
  
ax.set_xlabel('N-filter length')
ax.set_ylabel('D-filter length')
ax.set_zlabel('Q-factor (dB)')
ax.invert_xaxis()
ax.legend(loc=(-0.18, 0.4))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13288532408411013, 'fraction_numerical': 0.013440277476696292, 'mean_word_length': 2.1646090534979425, 'pattern_counts': {'":': 1, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 23, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '50932', 'license': 'MIT', 'n_tokens_mistral': 1593, 'n_tokens_neox': 1517, 'n_words': 331, 'path': 'docs/notebooks/benchmark_few_taps.ipynb', 'repo_name': 'remifan/gdbp_study', 'size': '4613'}