lczerolens#

Main module for the lczerolens package.

Submodules#

Classes#

LczeroBoard

A class for wrapping the LczeroBoard class.

LczeroModel

Class for wrapping the LCZero models.

Package Contents#

class lczerolens.LczeroBoard(fen=STARTING_FEN, *, chess960=False)[source]#

Bases: chess.Board

A class for wrapping the LczeroBoard class.

Parameters:

  • fen (Optional[str])

  • chess960 (bool)

static get_plane_order(us)[source]#

Get the plane order for the given us view.

Parameters:

us (bool) – The us_them tuple.

Returns:

The plane order.

Return type:

str

static get_piece_index(piece, us, plane_order=None)[source]#

Converts a piece to its index in the plane order.

Parameters:

  • piece (str) – The piece to convert.

  • us (bool) – The us_them tuple.

  • plane_order (Optional[str]) – The plane order.

Returns:

The index of the piece in the plane order.

Return type:

int

to_config_tensor(us=None)[source]#

Converts a LczeroBoard to a tensor based on the pieces configuration.

Parameters:

us (Optional[bool]) – The us_them tuple.

Returns:

The 13x8x8 tensor.

Return type:

torch.Tensor

to_input_tensor(*, input_encoding=InputEncoding.INPUT_CLASSICAL_112_PLANE)[source]#

Create the lc0 input tensor from the history of a game.

Parameters:

input_encoding (InputEncoding) – The input encoding method.

Returns:

The 112x8x8 tensor.

Return type:

torch.Tensor

static encode_move(move, us)[source]#

Converts a chess.Move object to an index.

Parameters:

  • move (chess.Move) – The chess move to encode.

  • us (bool) – The side to move (True for white, False for black).

Returns:

The encoded move index.

Return type:

int

decode_move(index)[source]#

Converts an index to a chess.Move object.

Parameters:

index (int) – The index to convert.

Returns:

The chess move.

Return type:

chess.Move

Gets the legal indices.

Returns:

Tensor containing indices of legal moves.

Return type:

torch.Tensor

Gets the next legal boards.

Parameters:

n_history (int, optional) – Number of previous positions to keep in the move stack, by default 7.

Returns:

Generator yielding board positions after each legal move.

Return type:

Generator[LczeroBoard, None, None]

render_heatmap(heatmap, square=None, vmin=None, vmax=None, arrows=None, normalise='none', save_to=None, cmap_name='RdYlBu_r', alpha=1.0, relative_board_view=True, heatmap_mode='relative_flip')[source]#

Render a heatmap on the board.

Parameters:

  • heatmap (torch.Tensor or numpy.ndarray) – The heatmap values to visualize on the board (64,) or (8, 8).

  • square (Optional[str], default=None) – Chess square to highlight (e.g. ‘e4’).

  • vmin (Optional[float], default=None) – Minimum value for the colormap normalization.

  • vmax (Optional[float], default=None) – Maximum value for the colormap normalization.

  • arrows (Optional[List[Tuple[str, str]]], default=None) – List of arrow tuples (from_square, to_square) to draw on board.

  • normalise (str, default="none") – Normalization method. Use “abs” for absolute value normalization.

  • save_to (Optional[str], default=None) – Path to save the visualization. If None, returns the figure.

  • cmap_name (str, default="RdYlBu_r") – Name of matplotlib colormap to use.

  • alpha (float, default=1.0) – Opacity of the heatmap overlay.

  • relative_board_view (bool, default=True) – Whether to use the relative board view.

  • heatmap_mode (str, default="relative_flip") – Use “relative_flip” if the heatmap corresponds to a relative flip of the board, “relative_rotation” if it corresponds to a relative rotation of the board, or “absolute” if it is already in the correct orientation.

Returns:

If save_to is None, returns (SVG string, matplotlib figure). If save_to is provided, saves files and returns None.

Return type:

Union[Tuple[str, matplotlib.figure.Figure], None]

Raises:

ValueError – If save_to is provided and does not end with .svg.

class lczerolens.LczeroModel(module, out_keys, **kwargs)[source]#

Bases: tensordict.nn.TensorDictModule

Class for wrapping the LCZero models.

Parameters:

  • module (torch.nn.Module)

  • out_keys (List[str])

prepare_boards(*boards, input_encoding=InputEncoding.INPUT_CLASSICAL_112_PLANE)[source]#

Prepares the boards for the model.

Parameters:

  • *boards (LczeroBoard) – The boards to prepare.

  • input_encoding (InputEncoding, optional) – The encoding of the boards.

Returns:

The prepared boards.

Return type:

torch.Tensor

forward(inputs, prepare_kwargs=None, **kwargs)[source]#
Parameters:

  • inputs (Union[TensorDict, Iterable[LczeroBoard], torch.Tensor]) – The inputs to the model.

  • prepare_kwargs (Optional[Dict[str, Any]], optional) – Keyword arguments to pass to the prepare_boards method, by default None

  • **kwargs (Any) – Additional keyword arguments to pass to the super().forward method.

Returns:

The output of the model.

Return type:

TensorDict

_call_module(tensors, **kwargs)[source]#
Parameters:

  • tensors (Sequence[torch.Tensor])

  • kwargs (Any)

Return type:

Sequence[torch.Tensor]

classmethod from_model(model, **kwargs)[source]#

Creates a wrapper from a model.

Parameters:

  • model (nn.Module) – The model to wrap.

  • **kwargs (Any) – Additional keyword arguments to pass to the super().__init__ method.

Returns:

The wrapped model instance

Return type:

LczeroModel

classmethod from_path(model_path, **kwargs)[source]#

Creates a wrapper from a model path.

Parameters:

model_path (str) – Path to the model file (.onnx or .pt)

Returns:

The wrapped model instance

Return type:

LczeroModel

Raises:

NotImplementedError – If the model file extension is not supported

classmethod from_onnx_path(onnx_model_path, check=True, **kwargs)[source]#

Builds a model from an ONNX file path.

Parameters:

  • onnx_model_path (str) – Path to the ONNX model file

  • check (bool, optional) – Whether to perform shape inference check, by default True

Returns:

The wrapped model instance

Return type:

LczeroModel

Raises:

  • FileNotFoundError – If the model file does not exist

  • ValueError – If the model could not be loaded

classmethod from_torch_path(torch_model_path, weights_only=False, **kwargs)[source]#

Builds a model from a PyTorch file path.

Parameters:

  • torch_model_path (str) – Path to the PyTorch model file

  • weights_only (bool)

Returns:

The wrapped model instance

Return type:

LczeroModel

Raises:

  • FileNotFoundError – If the model file does not exist

  • ValueError – If the model could not be loaded or is not a valid model type

push_to_hf(repo_id, create_if_not_exists=True, create_kwargs=None, path_in_repo='model.pt', **kwargs)[source]#

Pushes the model to the Hugging Face Hub.

Parameters:

  • repo_id (str) – The repository id to push the model to.

  • create_if_not_exists (bool, optional) – Whether to create the repository if it does not exist, by default True

  • create_kwargs (Optional[Dict[str, Any]], optional) – Additional keyword arguments to pass to the create_repo method.

  • path_in_repo (str, optional) – The path in the repository to save the model to.

  • **kwargs (Any) – Additional keyword arguments to pass to the upload_file method.

Raises:

ImportError – If the huggingface_hub library is not installed.

classmethod from_hf(repo_id, filename='model.pt', hf_hub_kwargs=None, **kwargs)[source]#

Loads a model from the Hugging Face Hub.

Parameters:

  • repo_id (str) – The repository id to load the model from.

  • filename (str) – The filename of the model to load.

  • hf_hub_kwargs (Optional[Dict[str, Any]], optional) – Additional keyword arguments to pass to the hf_hub_download method.

  • **kwargs (Any) – Additional keyword arguments to pass to the from_path method.

Returns:

The loaded model instance

Return type:

LczeroModel

Raises:

ImportError – If the huggingface_hub library is not installed.

static _get_output_names(model)[source]#

Returns the output names of the model.

Parameters:

model (nn.Module) – The model to get the output names from.

Returns:

The output names of the model.

Return type:

List[str]