114class MazeDataset(GPTDataset[MazeDatasetConfig]):  # noqa: PLW1641
115	"""a maze dataset class. This is a collection of solved mazes, and should be initialized via `MazeDataset.from_config`"""
116
117	def __init__(
118		self,
119		cfg: MazeDatasetConfig,
120		mazes: typing.Sequence[SolvedMaze],
121		generation_metadata_collected: dict | None = None,
122	) -> None:
123		"""initialize a maze dataset from a config and a list of solved mazes"""
124		super().__init__()
125		self.cfg: MazeDatasetConfig = cfg
126		self.mazes: list[SolvedMaze] = list(mazes)
127		self.generation_metadata_collected: dict | None = generation_metadata_collected
128
129	# TYPING: error: Return type "MazeDataset" of "from_config" incompatible with return type "T_Dataset" in supertype "GPTDataset"  [override]
130	@classmethod
131	def from_config(  # type: ignore[override]
132		cls,
133		# TYPING: error: Argument 1 of "from_config" is incompatible with supertype "GPTDataset"; supertype defines the argument type as "T_DatasetConfig"  [override]
134		cfg: MazeDatasetConfig,  # type: ignore[override]
135		do_generate: bool = True,
136		load_local: bool = True,
137		save_local: bool = True,
138		zanj: ZANJ | None = None,
139		do_download: bool = True,
140		local_base_path: Path = Path("data/maze_dataset"),
141		except_on_config_mismatch: bool = True,
142		allow_generation_metadata_filter_mismatch: bool = True,
143		verbose: bool = False,
144		**kwargs,
145	) -> "MazeDataset":
146		"""create a maze dataset from a config
147
148		priority of loading:
149		1. load from local
150		2. download
151		3. generate
152
153		"""
154		return cast(
155			"MazeDataset",
156			super().from_config(
157				cfg=cfg,
158				do_generate=do_generate,
159				load_local=load_local,
160				save_local=save_local,
161				zanj=zanj,
162				do_download=do_download,
163				local_base_path=local_base_path,
164				except_on_config_mismatch=except_on_config_mismatch,
165				allow_generation_metadata_filter_mismatch=allow_generation_metadata_filter_mismatch,
166				verbose=verbose,
167				**kwargs,
168			),
169		)
170
171	def data_hash(self) -> int:
172		"""return a hash of the data"""
173		return stable_hash(str(tuple([x.serialize() for x in self.mazes])))
174
175	def __getitem__(self, i: int) -> SolvedMaze:
176		"""get a maze by index"""
177		return self.mazes[i]
178
179	def __iter__(self) -> typing.Iterator[SolvedMaze]:
180		"""iterate over the mazes"""
181		return iter(self.mazes)
182
183	def __deepcopy__(self, memo) -> "MazeDataset":  # noqa: ANN001
184		"""deepcopy the dataset
185
186		FIX: this isnt actually a deepcopy I think?
187		"""
188		return MazeDataset.load(self._serialize_full())
189
190	# TYPING: get type hints on the tokenizer here
191	@overload
192	def as_tokens(
193		self,
194		maze_tokenizer,  # noqa: ANN001
195		limit: int | None = None,
196		join_tokens_individual_maze: Literal[False] = False,
197	) -> list[list[str]]: ...
198	@overload
199	def as_tokens(
200		self,
201		maze_tokenizer,  # noqa: ANN001
202		limit: int | None = None,
203		join_tokens_individual_maze: Literal[True] = True,
204	) -> list[str]: ...
205	def as_tokens(
206		self,
207		maze_tokenizer,  # TODO: MazeTokenizer
208		limit: int | None = None,
209		join_tokens_individual_maze: bool = False,
210	) -> list[list[str]] | list[str]:
211		"""return the dataset as tokens according to the passed `maze_tokenizer`
212
213		the `maze_tokenizer` should be either a `MazeTokenizer` or a `MazeTokenizerModular`
214
215		if `join_tokens_individual_maze` is True, then the tokens of each maze are
216		joined with a space, and the result is a list of strings.
217		i.e.:
218
219			>>> dataset.as_tokens(join_tokens_individual_maze=False)
220			[["a", "b", "c"], ["d", "e", "f"]]
221			>>> dataset.as_tokens(join_tokens_individual_maze=True)
222			["a b c", "d e f"]
223		"""
224		output: list[list[str]] = [
225			maze.as_tokens(maze_tokenizer) for maze in self.mazes[:limit]
226		]
227		if join_tokens_individual_maze:
228			return [" ".join(tokens) for tokens in output]
229		else:
230			return output
231
232	def __len__(self) -> int:
233		"""return the number of mazes in the dataset"""
234		return len(self.mazes)
235
236	def __eq__(self, other: object) -> bool:
237		"""compare two datasets"""
238		if not isinstance(other, MazeDataset):
239			raise NotImplementedError(
240				"can only compare with other MazeDataset objects",
241			)
242		# TODO: compare hashes of data instead of the data itself?
243		return self.cfg == other.cfg and self.mazes == other.mazes
244
245	def assert_equal(self, other: "MazeDataset") -> None:
246		"""assert that two datasets are equal"""
247		assert isinstance(other, MazeDataset)
248		assert self.cfg == other.cfg, f"{self.cfg.diff(other.cfg) = }"
249		assert self.mazes == other.mazes, f"{self.mazes = }, {other.mazes = }"
250
251	@classmethod
252	def generate(
253		cls,
254		cfg: MazeDatasetConfig,
255		gen_parallel: bool = False,
256		pool_kwargs: dict | None = None,
257		verbose: bool = False,
258		# TODO: what to do when unexpected kwargs are passed?
259		**kwargs,  # noqa: ARG003
260	) -> "MazeDataset":
261		"""Generate a maze dataset given a config and some generation parameters"""
262		# Copy the config to avoid modifying the original
263		cfg_cpy: MazeDatasetConfig = MazeDatasetConfig.load(
264			json.loads(json.dumps(cfg.serialize())),
265		)
266
267		if pool_kwargs is None:
268			pool_kwargs = dict()
269		maze_indexes: Int[np.ndarray, " maze_index"] = np.arange(cfg_cpy.n_mazes)  # type: ignore[assignment]
270
271		solved_mazes: list[SolvedMaze | None]
272		# Configure tqdm for progress bar
273		tqdm_kwargs: dict = dict(
274			total=cfg_cpy.n_mazes,
275			unit="maze",
276			desc="generating & solving mazes",
277			disable=not verbose,
278		)
279		# TODO: don't use the global unless generating in parallel!
280		if gen_parallel:
281			with multiprocessing.Pool(
282				**pool_kwargs,
283				initializer=_maze_gen_init_worker,
284				initargs=(cfg_cpy,),
285			) as pool:
286				solved_mazes = list(
287					tqdm.tqdm(
288						pool.imap(_generate_maze_helper, maze_indexes),
289						**tqdm_kwargs,
290					),
291				)
292
293		else:
294			_maze_gen_init_worker(cfg_cpy)
295			solved_mazes = list(
296				tqdm.tqdm(
297					map(
298						# TYPING:  error: Argument 1 to "map" has incompatible type "Callable[[int], SolvedMaze | None]"; expected "Callable[[str], SolvedMaze | None]"  [arg-type]
299						# why does it think tolist() returns a string?
300						_generate_maze_helper,  # type: ignore[arg-type]
301						maze_indexes.tolist(),
302					),
303					**tqdm_kwargs,
304				),
305			)
306
307		# Filter out None values explicitly after ensuring all results are collected
308		solved_mazes_: list[SolvedMaze] = [
309			maze for maze in solved_mazes if maze is not None
310		]
311		# solved_mazes_ = list(filter(lambda x: x is not None, solved_mazes))
312
313		# Update the config with the actual number of mazes
314		cfg_cpy.n_mazes = len(solved_mazes_)
315
316		dataset: MazeDataset = cls(
317			cfg=cfg_cpy,
318			mazes=solved_mazes_,
319		)
320
321		dataset.update_self_config()  # Call `update_self_config()` to ensure the dataset's config reflects changes
322
323		np.random.seed(cfg_cpy.seed)  # Reset the seed to the value in the config copy
324
325		return dataset
326
327	@classmethod
328	def download(cls, cfg: MazeDatasetConfig, **kwargs) -> "MazeDataset":
329		"(not implemented yet!) download a maze dataset from the internet"
330		raise NotImplementedError("not implemented yet")
331
332	@classmethod
333	def load(cls: "type[MazeDataset]", data: JSONdict) -> "MazeDataset":
334		"""load from zanj/json"""
335		if data[_FORMAT_KEY] == "MazeDataset:minimal":
336			return cls._load_minimal(data)
337		elif data[_FORMAT_KEY] == "MazeDataset:minimal_soln_cat":
338			return cls._load_minimal_soln_cat(data)
339		elif data[_FORMAT_KEY] == "MazeDataset":
340			if (
341				SERIALIZE_MINIMAL_THRESHOLD == -1
342			):  # Allow access to `_load_legacy` for profiling
343				return cls._load_legacy(data)
344			return cls._load_full(data)
345		else:
346			err_msg: str = f"`_FORMAT_KEY` string {data[_FORMAT_KEY] = } is not a recognized `MazeDataset` format. ({_FORMAT_KEY = })"
347			raise KeyError(
348				err_msg,
349			)
350
351	@classmethod
352	def _load_full(cls, data: JSONdict) -> "MazeDataset":
353		assert data[_FORMAT_KEY] == "MazeDataset"
354		return cls(
355			cfg=MazeDatasetConfig.load(data["cfg"]),  # type: ignore[arg-type]
356			mazes=load_item_recursive(data["mazes"], tuple()),
357			generation_metadata_collected=data["generation_metadata_collected"],  # type: ignore[arg-type]
358		)
359
360	@classmethod
361	def _load_minimal(cls, data: JSONdict) -> "MazeDataset":
362		assert data[_FORMAT_KEY] == "MazeDataset:minimal"
363		return cls(
364			cfg=MazeDatasetConfig.load(data["cfg"]),  # type: ignore[arg-type]
365			generation_metadata_collected=data["generation_metadata_collected"],  # type: ignore[arg-type]
366			mazes=[
367				SolvedMaze(
368					clist,
369					soln[:slen, ...],
370				)
371				for clist, slen, soln in zip(
372					load_item_recursive(data["maze_connection_lists"], tuple()),
373					load_item_recursive(data["maze_solution_lengths"], tuple()),
374					load_item_recursive(data["maze_solutions"], tuple()),
375					strict=False,
376					# load_item_recursive(data["maze_endpoints"], tuple()),
377				)
378			],
379		)
380
381	@classmethod
382	def _load_minimal_soln_cat(cls, data: JSONdict) -> "MazeDataset":
383		assert data[_FORMAT_KEY] == "MazeDataset:minimal_soln_cat"
384
385		maze_solution_lengths = load_item_recursive(
386			data["maze_solution_lengths"],
387			tuple(),
388		)
389		maze_solutions_concat = load_item_recursive(
390			data["maze_solutions_concat"],
391			tuple(),
392		)
393		maze_solutions = np.split(
394			maze_solutions_concat,
395			np.cumsum(maze_solution_lengths)[:-1],
396			axis=0,
397		)
398
399		return cls(
400			cfg=load_item_recursive(data["cfg"], tuple()),
401			generation_metadata_collected=load_item_recursive(
402				data["generation_metadata_collected"],
403				tuple(),
404			),
405			mazes=[
406				SolvedMaze(
407					connection_list=clist,
408					solution=soln,
409				)
410				for clist, soln in zip(
411					load_item_recursive(data["maze_connection_lists"], tuple()),
412					# load_item_recursive(data["maze_endpoints"], tuple()),
413					maze_solutions,
414					strict=False,
415				)
416			],
417		)
418
419	@classmethod
420	def _load_legacy(cls, data: JSONdict) -> "MazeDataset":
421		"""Legacy `load` method from <0.5.2. Used exclusively for profiling comparison."""
422		assert data[_FORMAT_KEY] == "MazeDataset"
423		return cls(
424			**{
425				key: load_item_recursive(data[key], tuple())
426				for key in ["cfg", "mazes", "generation_metadata_collected"]
427			},
428		)
429
430	def serialize(self) -> JSONdict:
431		"""serialize to zanj/json"""
432		if (
433			SERIALIZE_MINIMAL_THRESHOLD is not None
434			and len(self) >= SERIALIZE_MINIMAL_THRESHOLD
435		):
436			return self._serialize_minimal()
437		return self._serialize_full()
438
439	def _serialize_full(self) -> JSONdict:
440		return {
441			_FORMAT_KEY: "MazeDataset",
442			"cfg": json_serialize(self.cfg),
443			"fname": self.cfg.to_fname(),
444			"mazes": json_serialize(self.mazes),
445			"generation_metadata_collected": json_serialize(
446				self.generation_metadata_collected,
447			),
448		}
449
450	def _serialize_minimal(self) -> JSONdict:
451		"alternate serialization where metadata is collected and mazes are stored in concatenated form"
452		filtered_meta: MazeDataset
453		if self.generation_metadata_collected is None:
454			filtered_meta = self.filter_by.collect_generation_meta()
455		else:
456			filtered_meta = self
457
458		max_solution_len: int = max(m.solution.shape[0] for m in filtered_meta.mazes)
459		n_mazes: int = len(filtered_meta.mazes)
460		grid_n: int = filtered_meta.cfg.grid_n
461
462		maze_connection_lists: np.ndarray = np.empty(
463			(n_mazes, 2, grid_n, grid_n),
464			dtype=np.bool_,
465		)
466		# maze_endpoints: np.ndarray = np.empty((n_mazes, 2, 2), dtype=np.int8)
467		maze_solution_lengths: np.ndarray = np.empty((n_mazes,), dtype=np.int32)
468		maze_solutions: np.ndarray = np.empty(
469			(n_mazes, max_solution_len, 2),
470			dtype=np.int8,
471		)
472
473		for idx, maze in enumerate(filtered_meta.mazes):
474			maze_connection_lists[idx] = maze.connection_list
475			# maze_endpoints[idx] = np.array([maze.start_pos, maze.end_pos])
476			maze_solution_lengths[idx] = maze.solution.shape[0]
477			maze_solutions[idx, : maze.solution.shape[0]] = maze.solution
478
479		return {
480			_FORMAT_KEY: "MazeDataset:minimal",
481			"cfg": json_serialize(filtered_meta.cfg),
482			"fname": filtered_meta.cfg.to_fname(),
483			"generation_metadata_collected": json_serialize(
484				filtered_meta.generation_metadata_collected,
485			),
486			"maze_connection_lists": maze_connection_lists,  # type: ignore[dict-item]
487			# "maze_endpoints": maze_endpoints,
488			"maze_solution_lengths": maze_solution_lengths,  # type: ignore[dict-item]
489			"maze_solutions": maze_solutions,  # type: ignore[dict-item]
490		}
491
492	def _serialize_minimal_soln_cat(self: "MazeDataset") -> JSONdict:
493		"alternate serialization where metadata is collected, and mazes and their solutions are stored in concatenated form"
494		filtered_meta: MazeDataset
495		if self.generation_metadata_collected is None:
496			filtered_meta = self.filter_by.collect_generation_meta()
497		else:
498			filtered_meta = self
499
500		maze_solution_lengths: np.ndarray = np.array(
501			[m.solution.shape[0] for m in filtered_meta.mazes],
502			dtype=np.int32,
503		)
504		n_mazes: int = len(filtered_meta.mazes)
505		grid_n: int = filtered_meta.cfg.grid_n
506		total_solution_len: int = np.sum(maze_solution_lengths)
507
508		maze_connection_lists: np.ndarray = np.empty(
509			(n_mazes, 2, grid_n, grid_n),
510			dtype=np.bool_,
511		)
512		maze_endpoints: np.ndarray = np.empty((n_mazes, 2, 2), dtype=np.int8)
513		maze_solutions_concat: np.ndarray = np.empty(
514			(total_solution_len, 2),
515			dtype=np.int8,
516		)
517
518		solutions_running_idx: int = 0
519		for idx, maze in enumerate(filtered_meta.mazes):
520			maze_connection_lists[idx] = maze.connection_list
521			maze_endpoints[idx] = np.array([maze.start_pos, maze.end_pos])
522			soln_len: int = maze.solution.shape[0]
523			maze_solution_lengths[idx] = soln_len
524			maze_solutions_concat[
525				solutions_running_idx : solutions_running_idx + soln_len
526			] = maze.solution
527			solutions_running_idx += soln_len
528
529		return {
530			_FORMAT_KEY: "MazeDataset:minimal_soln_cat",
531			"cfg": json_serialize(filtered_meta.cfg),
532			"fname": filtered_meta.cfg.to_fname(),
533			"generation_metadata_collected": json_serialize(
534				filtered_meta.generation_metadata_collected,
535			),
536			"maze_connection_lists": maze_connection_lists,  # type: ignore[dict-item]
537			"maze_endpoints": maze_endpoints,  # type: ignore[dict-item]
538			"maze_solution_lengths": maze_solution_lengths,  # type: ignore[dict-item]
539			"maze_solutions_concat": maze_solutions_concat,  # type: ignore[dict-item]
540		}
541
542	def update_self_config(self) -> None:
543		"""update the config to match the current state of the dataset (number of mazes, such as after filtering)"""
544		if self.cfg.n_mazes != len(self.mazes):
545			warnings.warn(
546				f"updating config n_mazes from {self.cfg.n_mazes} to {len(self.mazes)}",
547			)
548			self.cfg.n_mazes = len(self.mazes)
549
550	def custom_maze_filter(
551		self,
552		method: typing.Callable[[SolvedMaze], bool],
553		**kwargs,
554	) -> "MazeDataset":
555		"""filter the dataset using a custom method"""
556		output: MazeDataset = MazeDataset(
557			cfg=copy.deepcopy(self.cfg),
558			mazes=[m for m in self.mazes if method(m, **kwargs)],
559		)
560		output.cfg.applied_filters.append(
561			{
562				"name": f"__custom__:{method.__name__}",
563				"kwargs": kwargs,
564			},
565		)
566		output.update_self_config()
567		return output

257@serializable_dataclass(kw_only=True, methods_no_override=["serialize"])
258class MazeDatasetConfig(MazeDatasetConfig_base):  # type: ignore[misc]
259	"""config object which is passed to `MazeDataset.from_config` to generate or load a dataset
260
261	# Parameters:
262	- `name : str`
263		name of the dataset -- this can be anything, but should be filesystem safe since we use it in the `fname`
264	- `grid_n : int`
265		grid size of the maze (number of rows/columns)
266	- `n_mazes : int`
267		number of mazes to request. For some combinations of `endpoint_kwargs` and `maze_ctor`, not all mazes might successfully generate.
268		see `EndpointKwargsType` for more details.
269	- `maze_ctor : Callable`
270		maze generator function. This should be a function that takes a grid size and returns a maze.
271		This will usually be one of the functions in `LatticeMazeGenerators`.
272	- `maze_ctor_kwargs : dict`
273		keyword arguments to pass to the maze generator function. Specific to the `maze_ctor` you are using.
274	- `endpoint_kwargs : EndpointKwargsType`
275		keyword arguments passed to `LatticeMaze.generate_random_path()`. see `EndpointKwargsType` for more info.
276	- `applied_filters : list[dict]`
277		list of filters that have been applied to the dataset. We recommend applying filters to datasets directly,
278		but these are stored with the config in case you want to re-generate the dataset with the same filters.
279
280	"""
281
282	@property
283	def config_version(self) -> str:
284		"""return the version of the config. added in maze_dataset v1.3.0, previous versions had no dataset config"""
285		return "1.0"
286
287	@property
288	def versions(self) -> dict:
289		"""return the versions of the config and the maze_dataset"""
290		return dict(
291			config=self.config_version,
292			maze_dataset=importlib.metadata.version("maze_dataset"),
293		)
294
295	def serialize(self) -> dict:
296		"serialize the MazeDatasetConfig with all fields and fname"
297		return {
298			**self._serialize_base(
299				applied_filters__skip__collect_generation_meta=False
300			),
301			"fname": self.to_fname(),
302			"versions": self.versions,
303		}
304
305	def summary(self) -> dict:
306		"""return a summary of the config"""
307		# do we run this to make sure it doesn't error?
308		super_summary: dict = super().summary()
309		assert super_summary
310		self_ser: dict = self.serialize()
311		return dict(
312			name=self.name,
313			fname=self.to_fname(),
314			sdc_hash=self.stable_hash_cfg(),
315			seed=self.seed,
316			seq_len_min=self.seq_len_min,
317			seq_len_max=self.seq_len_max,
318			applied_filters=self.applied_filters,
319			grid_n=self_ser["grid_n"],
320			n_mazes=self_ser["n_mazes"],
321			maze_ctor_name=self_ser["maze_ctor"]["__name__"],
322			maze_ctor_kwargs=self_ser["maze_ctor_kwargs"],
323			endpoint_kwargs=self_ser["endpoint_kwargs"],
324		)
325
326	def _to_ps_array(self) -> _PercolationSuccessArray:
327		"""Convert this config to a [p, grid_n, deadends, endpoints_not_equal, generator_func] vector.
328
329		used in predicting the success rate
330		"""
331		try:
332			assert self.maze_ctor.__name__ in _GENERATORS_PERCOLATED, (
333				f"generator not supported, must be a percolation generator\n{self.maze_ctor.__name__ = }, {_GENERATORS_PERCOLATED = }"
334			)
335			assert "p" in self.maze_ctor_kwargs, (
336				f"maze_ctor_kwargs must have a 'p' (percolation value) key: {self.maze_ctor_kwargs = }"
337			)
338			assert not self.endpoint_kwargs.get("except_on_no_valid_endpoint", True), (
339				f"except_on_no_valid_endpoint must be False, or else if any maze fails to generate, the whole dataset will fail: {self.endpoint_kwargs = }"
340			)
341		except AssertionError as e:
342			err_msg: str = f"invalid config for percolation success prediction: {self.summary() = }"
343			raise NoPercolationInConfigError(
344				err_msg,
345			) from e
346
347		endpoints_unique_flag: int = int(
348			# we are pretty sure it will be an int or bool here
349			self.endpoint_kwargs.get("endpoints_not_equal", True),  # type: ignore[arg-type]
350		)
351
352		# adjustment for bknutson0
353		if not (
354			self.endpoint_kwargs.get("deadend_start", False)
355			and self.endpoint_kwargs.get("deadend_end", False)
356		):
357			# we didnt train on this, but if either endpoint is not required to be in a dead end
358			# then  requiring the endpoints to be unique does not really affect the success rate
359			# (except for very small percolation values, pure percolation generation)
360			endpoints_unique_flag = 0
361
362		return np.array(
363			[
364				float(self.maze_ctor_kwargs["p"]),
365				float(self.grid_n),
366				float(
367					int(
368						self.endpoint_kwargs.get("deadend_start", False)  # type: ignore[arg-type]
369						or self.endpoint_kwargs.get("deadend_end", False),
370					),
371				),
372				float(endpoints_unique_flag),
373				float(_GENERATORS_PERCOLATED.index(self.maze_ctor.__name__)),
374			],
375			dtype=np.float64,
376		)
377
378	@classmethod
379	def _from_ps_array(
380		cls,
381		arr: _PercolationSuccessArray,
382		name: str = "predict",
383		n_mazes: int = 100,
384		**kwargs,
385	) -> "MazeDatasetConfig":
386		"""Reconstruct a config from an array [p, grid_n, deadends, endpoints_not_equal, generator_func] and other config parameters.
387
388		# Returns:
389		- `MazeDatasetConfig`
390			Config corresponding to `arr`
391		"""
392		return cls(
393			name=name,
394			grid_n=int(arr[1]),
395			n_mazes=n_mazes,
396			maze_ctor=GENERATORS_MAP[_GENERATORS_PERCOLATED[int(arr[4])]],
397			maze_ctor_kwargs={"p": float(arr[0])},
398			endpoint_kwargs=dict(
399				deadend_start=bool(arr[2]),
400				deadend_end=bool(arr[2]),
401				endpoints_not_equal=bool(arr[3]),
402				except_on_no_valid_endpoint=False,
403			),
404			**kwargs,
405		)
406
407	def success_fraction_estimate(
408		self,
409		except_if_all_success_expected: bool = False,
410	) -> float:
411		"""Estimate the success fraction of this config.
412
413		only valid when the generator is a percolation generator,
414		and endpoints are enforced to be dead ends
415
416		more information on where this comes from can be found in
417		- `cfg_success_predict_fn()` from `maze_dataset.dataset.success_predict_math`
418		- `estimate_dataset_fractions.ipynb`
419		- `maze_dataset.benchmarks.sweep_fit`
420
421		# Parameters:
422		- `except_if_all_success_expected : bool`
423			if `True`, don't raise an error if the success fraction is below the threshold.
424			will always return `1.0` if the config is not expected to fail
425
426		# Returns:
427		- `float`
428			estimated success fraction
429
430		# Raises:
431		- `NoPercolationInConfigError` : if the config is not expected to fail, and `except_if_all_success_expected` is `False`
432		"""
433		try:
434			return cfg_success_predict_fn(self)
435
436		except NoPercolationInConfigError as e:
437			if except_if_all_success_expected:
438				raise e  # noqa: TRY201
439			return 1.0
440
441	def success_fraction_compensate(
442		self,
443		safety_margin: float = 1.2,
444		except_if_all_success_expected: bool = False,
445		epsilon: float = 1e-2,
446	) -> "MazeDatasetConfig":
447		"""return a new `MazeDatasetConfig` like this one with `n_mazes` adjusted to compensate for the success fraction
448
449		calls `MazeDatasetConfig.success_fraction_estimate()` to get the success fraction, and then
450		computes the new number of mazes as `n_mazes = n_mazes * safety_margin / success_fraction + 1`
451
452		more information on where this comes from can be found in
453		- `cfg_success_predict_fn()` from `maze_dataset.dataset.success_predict_math`
454		- `estimate_dataset_fractions.ipynb`
455		- `maze_dataset.benchmarks.sweep_fit`
456
457		# Parameters:
458		- `safety_margin : float`
459			safety margin to apply to the success fraction estimate
460			(defaults to `1.2`, or 20% more mazes than estimated)
461		- `except_if_all_success_expected : bool`
462			if `True`, don't raise an error if the success fraction is below the threshold.
463			this is passed to `MazeDatasetConfig.success_fraction_estimate`.
464			if your config isn't expected to fail, passing this might mean you generate more mazes than needed
465			since `safety_margin` is still applied.
466			(defaults to `False`)
467		- `epsilon : float`
468			raise `SuccessChanceTooSmallError` if the success fraction is below this threshold
469			(defaults to `1e-2`)
470
471		# Returns:
472		- `MazeDatasetConfig`
473			new config with adjusted `n_mazes`
474
475		# Raises:
476		- `SuccessChanceTooSmallError` : if the computed success fraction is below `epsilon`
477		"""
478		# compute and check the success fraction
479		success_fraction: float = self.success_fraction_estimate(
480			except_if_all_success_expected=except_if_all_success_expected,
481		)
482		if success_fraction < epsilon:
483			err_msg: str = (
484				f"{success_fraction = } is below the threshold of {epsilon = }"
485			)
486			raise SuccessChanceTooSmallError(
487				err_msg,
488			)
489
490		# compute the new number of mazes
491		n_mazes: int = self.n_mazes
492		new_n_mazes: int = int((n_mazes * safety_margin) / success_fraction) + 1
493
494		# put it in a new config and return
495		cfg_dict: dict = self.serialize()
496		cfg_dict["n_mazes"] = new_n_mazes
497		return MazeDatasetConfig.load(cfg_dict)