How to use xenonpy - 10 common examples

def test_transform_4x1_ravel_2(data):
    bc = BoxCox()
    bc.fit(data[0].ravel())
    trans = bc.transform((data[0].ravel())[:2])
    assert trans.shape == (data[2].ravel())[:2].shape
    assert np.all(trans == (data[2].ravel())[:2])
    inverse = bc.inverse_transform(trans)
    assert inverse.shape == (data[0].ravel())[:2].shape
    assert np.allclose(inverse, (data[0].ravel())[:2])

def test_trainer_fit_1(data):
    model = deepcopy(data[0])
    trainer = Trainer(model=model, optimizer=Adam(), loss_func=MSELoss())
    trainer.fit(*data[1])
    assert trainer.total_iterations == 200
    assert trainer.total_epochs == 200

    trainer.fit(*data[1], epochs=20)
    assert trainer.total_iterations == 220
    assert trainer.total_epochs == 220

    trainer.reset()
    assert trainer.total_iterations == 0
    assert trainer.total_epochs == 0

    trainer.fit(*data[1], epochs=20)
    assert trainer.total_iterations == 20
    assert trainer.total_epochs == 20

def test_trainer_prediction_1(data):
    model = deepcopy(data[0])
    trainer = Trainer(model=model, optimizer=Adam(lr=0.1), loss_func=MSELoss(), epochs=200)
    trainer.extend(TensorConverter())
    trainer.fit(*data[1], *data[1])

    trainer = Trainer(model=model).extend(TensorConverter())
    y_p = trainer.predict(data[1][0])
    assert np.any(np.not_equal(y_p, data[1][1].numpy()))
    assert np.allclose(y_p, data[1][1].numpy(), rtol=0, atol=0.2)

    y_p, y_t = trainer.predict(*data[1])
    assert np.any(np.not_equal(y_p, y_t))
    assert np.allclose(y_p, y_t, rtol=0, atol=0.2)

    val_set = DataLoader(TensorDataset(*data[1]), batch_size=50)
    y_p, y_t = trainer.predict(dataset=val_set)
    assert np.any(np.not_equal(y_p, y_t))
    assert np.allclose(y_p, y_t, rtol=0, atol=0.2)

def test_trainer_2(data):
    trainer = Trainer()
    with pytest.raises(RuntimeError, match='no model for training'):
        trainer.fit(*data[1])

    with pytest.raises(TypeError, match='parameter `m` must be a instance of '):
        trainer.model = {}

    trainer.model = data[0]
    assert isinstance(trainer.model, torch.nn.Module)
    with pytest.raises(RuntimeError, match='no loss function for training'):
        trainer.fit(*data[1])

    trainer.loss_func = MSELoss()
    assert trainer.loss_type == 'train_mse_loss'
    assert trainer.loss_func.__class__ == MSELoss
    with pytest.raises(RuntimeError, match='no optimizer for training'):
        trainer.fit(*data[1])

def test_trainer_prediction_1(data):
    model = deepcopy(data[0])
    trainer = Trainer(model=model, optimizer=Adam(lr=0.1), loss_func=MSELoss(), epochs=200)
    trainer.extend(TensorConverter())
    trainer.fit(*data[1], *data[1])

    trainer = Trainer(model=model).extend(TensorConverter())
    y_p = trainer.predict(data[1][0])
    assert np.any(np.not_equal(y_p, data[1][1].numpy()))
    assert np.allclose(y_p, data[1][1].numpy(), rtol=0, atol=0.2)

    y_p, y_t = trainer.predict(*data[1])
    assert np.any(np.not_equal(y_p, y_t))
    assert np.allclose(y_p, y_t, rtol=0, atol=0.2)

    val_set = DataLoader(TensorDataset(*data[1]), batch_size=50)
    y_p, y_t = trainer.predict(dataset=val_set)
    assert np.any(np.not_equal(y_p, y_t))
    assert np.allclose(y_p, y_t, rtol=0, atol=0.2)

    with pytest.raises(RuntimeError, match='parameters  and  are mutually exclusive'):
        trainer.predict(*data[1], dataset='not none')

def test_trainer_prediction_2():
    model = _Net(n_feature=2, n_hidden=10, n_output=2)

    n_data = np.ones((100, 2))
    x0 = np.random.normal(2 * n_data, 1)
    y0 = np.zeros(100)
    x1 = np.random.normal(-2 * n_data, 1)
    y1 = np.ones(100)

    x = np.vstack((x0, x1))
    y = np.concatenate((y0, y1))
    s = np.arange(x.shape[0])
    np.random.shuffle(s)
    x, y = x[s], y[s]

    trainer = Trainer(model=model, optimizer=Adam(lr=0.1), loss_func=CrossEntropyLoss(), epochs=200)
    trainer.extend(TensorConverter(x_dtype=torch.float32, y_dtype=torch.long, argmax=True))
    trainer.fit(x, y)

    y_p, y_t = trainer.predict(x, y)
    assert y_p.shape == (200,)
    assert np.all(y_p == y_t)

    # trainer.reset()
    val_set = DataLoader(ArrayDataset(x, y, dtypes=(torch.float, torch.long)), batch_size=20)
    trainer.extend(TensorConverter(x_dtype=torch.float32, y_dtype=torch.long, auto_reshape=False))
    y_p, y_t = trainer.predict(dataset=val_set)
    assert y_p.shape == (200, 2)

    y_p = np.argmax(y_p, 1)
    assert np.all(y_p == y_t)

def test_trainer_fit_2(data):
    model = deepcopy(data[0])
    trainer = Trainer(model=model, optimizer=Adam(), loss_func=MSELoss(), epochs=20)
    trainer.fit(*data[1], *data[1])
    assert trainer.total_iterations == 20
    assert trainer.total_epochs == 20
    assert (trainer.x_val, trainer.y_val) == data[1]

    train_set = DataLoader(TensorDataset(*data[1]))
    val_set = DataLoader(TensorDataset(*data[1]))
    trainer.fit(training_dataset=train_set, validation_dataset=val_set)
    assert trainer.total_iterations == 2020
    assert trainer.total_epochs == 40
    assert isinstance(trainer.validate_dataset, DataLoader)

with pytest.raises(TypeError, match='parameter `m` must be a instance of '):
        trainer.model = {}

    trainer.model = data[0]
    assert isinstance(trainer.model, torch.nn.Module)
    with pytest.raises(RuntimeError, match='no loss function for training'):
        trainer.fit(*data[1])

    trainer.loss_func = MSELoss()
    assert trainer.loss_type == 'train_mse_loss'
    assert trainer.loss_func.__class__ == MSELoss
    with pytest.raises(RuntimeError, match='no optimizer for training'):
        trainer.fit(*data[1])

    trainer.optimizer = Adam()
    assert isinstance(trainer.optimizer, torch.optim.Adam)
    assert isinstance(trainer._optimizer_state, dict)
    assert isinstance(trainer._init_states, dict)

    trainer.lr_scheduler = ExponentialLR(gamma=0.99)
    assert isinstance(trainer.lr_scheduler, torch.optim.lr_scheduler.ExponentialLR)

def test_trainer_fit_3(data):
    model = deepcopy(data[0])
    trainer = Trainer(model=model, optimizer=Adam(), loss_func=MSELoss(), epochs=5)
    trainer.fit(*data[1])
    assert len(trainer.checkpoints.keys()) == 0

    trainer.reset()
    assert trainer.total_iterations == 0
    assert trainer.total_epochs == 0
    assert len(trainer.get_checkpoint()) == 0

    trainer.fit(*data[1], checkpoint=True)
    assert len(trainer.get_checkpoint()) == 5
    assert isinstance(trainer.get_checkpoint(2), trainer.checkpoint_tuple)
    assert isinstance(trainer.get_checkpoint('cp_2'), trainer.checkpoint_tuple)

    with pytest.raises(TypeError, match='parameter  must be str or int'):
        trainer.get_checkpoint([])

def test_persist_1(data):
    model = deepcopy(data[0])
    trainer = Trainer(model=model, optimizer=Adam(lr=0.1), loss_func=MSELoss(), epochs=200)
    trainer.extend(TensorConverter(), Persist('model_dir'))
    trainer.fit(*data[1], *data[1])

    persist = trainer['persist']
    checker = persist._checker
    assert isinstance(persist, Persist)
    assert isinstance(checker.model, torch.nn.Module)
    assert isinstance(checker.describe, dict)
    assert isinstance(checker.files, list)
    assert set(checker.files) == {'model', 'init_state', 'model_structure', 'describe', 'training_info', 'final_state'}

    trainer = Trainer.load(checker)
    assert isinstance(trainer.training_info, pd.DataFrame)
    assert isinstance(trainer.model, torch.nn.Module)
    assert isinstance(trainer._training_info, list)
    assert trainer.optimizer is None