-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
68 lines (57 loc) · 2.38 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import torch
from torch.utils.data import Subset, TensorDataset
from typing import Optional, Union
TENSOR_INDICES = {
"X": 0,
"A": 1,
"T": 2,
"Y_obs": 3,
"Y": 4,
"prop_score": 5,
}
def _query_tensor_dataset(dataset: Union[TensorDataset, Subset], key: str) -> torch.Tensor:
"""Accessor for TensorDataset for the dataset format required by DCEM.
Args:
dataset (TensorDataset): A PyTorch TensorDataset object.
key (str): The name of the variable type (see keys of `TENSOR_INDICES`).
Returns:
Tensor: a Tensor containing all values for the specified variable.
"""
if isinstance(dataset, Subset):
return torch.cat([dataset.dataset.tensors[TENSOR_INDICES[key]][idx] for idx in dataset.indices], dim=0)
else:
return dataset.tensors[TENSOR_INDICES[key]]
def _enforce_bounds(
var,
varname: str,
_min: Optional[float] = float('-inf'),
_max: Optional[float] = float('inf'),
) -> None:
"""Checks if an array-like is in-bounds.
Args:
var (Any): Array-like variable; e.g., a tensor.
varname (str): Variable name (plaintext) to be used in the error message.
_min (Optional[float], optional): Minimum allowed value. Defaults to float('-inf').
_max (Optional[float], optional): Maximum allowed value. Defaults to float('inf').
Raises:
ValueError: Raises an error if the array is not completely within the specified bounds.
"""
if var.min() < _min or var.max() > _max:
raise ValueError(f"{varname} is not in [{_min}, {_max}]. Got min: {var.min()} and max {var.max()}")
def _enforce_scalar_bounds(
var,
varname: str,
_min: Optional[float] = float('-inf'),
_max: Optional[float] = float('inf'),
) -> None:
"""Same as _enforce_bounds, but for checking scalars.
Args:
var (Any): Scalar-like variable; e.g., a float.
varname (str): Variable name (plaintext) to be used in the error message.
_min (Optional[float], optional): Minimum allowed value. Defaults to float('-inf').
_max (Optional[float], optional): Maximum allowed value. Defaults to float('inf').
Raises:
ValueError: Raises an error if the value is not completely within the specified bounds.
"""
if var < _min or var > _max:
raise ValueError(f"{varname} is not in [{_min}, {_max}]. Got: {var}")