import pandas as pd

import larch as lx
from larch.util.namespaces import Namespace
spec_content = """expression,DA,SR,Walk,Bike,Transit
AUTO_TIME,InVehTime,InVehTime,,,
AUTO_COST,Cost,,,,
AUTO_COST*0.5,,Cost,,,
1,,ASC_SR,ASC_Walk,ASC_Bike,ASC_Transit
log(INCOME),,LogIncome_SR,LogIncome_Walk,LogIncome_Bike,LogIncome_Transit
WALK_TIME,,,NonMotorTime,,
BIKE_TIME,,,,NonMotorTime,
TRANSIT_IVTT,,,,,InVehTime
TRANSIT_OVTT,,,,,OutVehTime
TRANSIT_FARE,,,,,Cost
"""

from io import StringIO

spec = pd.read_csv(StringIO(spec_content)).set_index("expression")
spec
DA SR Walk Bike Transit
expression
AUTO_TIME InVehTime InVehTime NaN NaN NaN
AUTO_COST Cost NaN NaN NaN NaN
AUTO_COST*0.5 NaN Cost NaN NaN NaN
1 NaN ASC_SR ASC_Walk ASC_Bike ASC_Transit
log(INCOME) NaN LogIncome_SR LogIncome_Walk LogIncome_Bike LogIncome_Transit
WALK_TIME NaN NaN NonMotorTime NaN NaN
BIKE_TIME NaN NaN NaN NonMotorTime NaN
TRANSIT_IVTT NaN NaN NaN NaN InVehTime
TRANSIT_OVTT NaN NaN NaN NaN OutVehTime
TRANSIT_FARE NaN NaN NaN NaN Cost
hh, pp, tour, skims = lx.example(200, ["hh", "pp", "tour", "skims"])
Mode = Namespace(
    DA=1,
    SR=2,
    Walk=3,
    Bike=4,
    Transit=5,
)
# tour_dataset = lx.Dataset.construct.from_idco(tour.set_index('TOURID'), alts=Mode)

tour_dataset = tour.set_index("TOURID")
od_skims = lx.Dataset.construct.from_omx(skims)
dt = lx.DataTree(
    tour=tour_dataset,
    hh=hh.set_index("HHID"),
    person=pp.set_index("PERSONID"),
    od=od_skims,
    do=od_skims,
    relationships=(
        "tours.HHID @ hh.HHID",
        "tours.PERSONID @ person.PERSONID",
        "hh.HOMETAZ @ od.otaz",
        "tours.DTAZ @ od.dtaz",
        "hh.HOMETAZ @ do.dtaz",
        "tours.DTAZ @ do.otaz",
    ),
    root_node_name="tour",
)
spec.columns
Index(['DA', 'SR', 'Walk', 'Bike', 'Transit'], dtype='object')
dt.set_altnames(spec.columns)
dt.alts_name_to_id()
{'DA': np.int64(1),
 'SR': np.int64(2),
 'Walk': np.int64(3),
 'Bike': np.int64(4),
 'Transit': np.int64(5)}
import larch.model.from_spec
m = lx.Model.from_spec(spec, dt)
m.choice_co_code = "TOURMODE"
m.should_preload_data()
m.maximize_loglike()
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Larch Model Dashboard                                                                 ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ optimization complete                                                                 │
│ Log Likelihood Current =     -11021.350586 Best =     -11021.350586                   │
│ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓ │
│ ┃ Parameter                       Value           Gradient        Best           ┃ │
│ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩ │
│ │ ASC_Bike                       │  0.70431783    │ -0.11209822    │  0.70431785    │ │
│ │ ASC_SR                         │  3.1856348     │ -0.096607745   │  3.1856348     │ │
│ │ ASC_Transit                    │  2.8075233     │  0.47568321    │  2.8075233     │ │
│ │ ASC_Walk                       │  6.7836366     │  0.090206146   │  6.7836366     │ │
│ │ Cost                           │  0.5776521     │  1.2096558     │  0.5776521     │ │
│ │ InVehTime                      │ -0.20300829    │ -2.9287109     │ -0.20300829    │ │
│ │ LogIncome_Bike                 │ -0.30476749    │ -1.0374222     │ -0.30476749    │ │
│ │ LogIncome_SR                   │ -0.43725747    │ -1.375         │ -0.43725747    │ │
│ │ LogIncome_Transit              │ -0.37439233    │  4.9035034     │ -0.37439233    │ │
│ │ LogIncome_Walk                 │ -0.36716177    │  1.3070831     │ -0.36716177    │ │
│ │ NonMotorTime                   │ -0.24656629    │ -0.79376984    │ -0.24656629    │ │
│ │ OutVehTime                     │ -0.27982425    │  3.0588913     │ -0.27982425    │ │
│ └────────────────────────────────┴────────────────┴────────────────┴────────────────┘ │
└───────────────────────────────────────────────────────────────────────────────────────┘
      message: Optimization terminated successfully
      success: True
       status: 0
            x: [ 7.043e-01  3.186e+00  2.808e+00  6.784e+00  5.777e-01
                -2.030e-01 -3.048e-01 -4.373e-01 -3.744e-01 -3.672e-01
                -2.466e-01 -2.798e-01]
          nit: 72
          jac: [-1.121e-01 -9.661e-02  4.757e-01  9.021e-02  1.210e+00
                -2.929e+00 -1.037e+00 -1.375e+00  4.904e+00  1.307e+00
                -7.938e-01  3.059e+00]
         nfev: 173
         njev: 72
      n_cases: 20739
 total_weight: 20739.0
      logloss: 0.5314311483647958
      loglike: -11021.3505859375
dt.n_cases
20739
m.logsums()
array([ 0.09591921, -1.95402485, -1.94675511, ..., -0.4469249 ,
       -0.56657627, -0.13999554], shape=(20739,))
m.pvals = {"Cost": -0.1, "InVehTime": -0.2, "ASC_Bike": 2.0}
m.pf
value best initvalue minimum maximum nullvalue holdfast
param_name
ASC_Bike 2.000000 0.704318 0.0 -inf inf 0.0 0
ASC_SR 3.185635 3.185635 0.0 -inf inf 0.0 0
ASC_Transit 2.807523 2.807523 0.0 -inf inf 0.0 0
ASC_Walk 6.783637 6.783637 0.0 -inf inf 0.0 0
Cost -0.100000 0.577652 0.0 -inf inf 0.0 0
InVehTime -0.200000 -0.203008 0.0 -inf inf 0.0 0
LogIncome_Bike -0.304767 -0.304767 0.0 -inf inf 0.0 0
LogIncome_SR -0.437257 -0.437257 0.0 -inf inf 0.0 0
LogIncome_Transit -0.374392 -0.374392 0.0 -inf inf 0.0 0
LogIncome_Walk -0.367162 -0.367162 0.0 -inf inf 0.0 0
NonMotorTime -0.246566 -0.246566 0.0 -inf inf 0.0 0
OutVehTime -0.279824 -0.279824 0.0 -inf inf 0.0 0