nlsq.config module¶

Configuration management and environment setup.

Central configuration management for NLSQ package.

class nlsq.config.MemoryConfig(memory_limit_gb=8.0, gpu_memory_fraction=None, chunk_size_mb=None, out_of_memory_strategy='fallback', safety_factor=0.8, auto_chunk_threshold_gb=4.0, progress_reporting=True, min_chunk_size=1000, max_chunk_size=1000000)[source]

Bases: object

Configuration for memory management and GPU settings.

memory_limit_gb

Maximum memory limit in GB (default: 8.0)

Type:: float

gpu_memory_fraction

Fraction of GPU memory to use (0.0-1.0, None for automatic)

Type:: float | None

chunk_size_mb

Default chunk size in MB for data processing

Type:: int | None

out_of_memory_strategy

Strategy when out of memory: ‘fallback’, ‘reduce’, ‘error’

Type:: str

safety_factor

Safety factor for memory calculations (0.0-1.0)

Type:: float

auto_chunk_threshold_gb

Automatically enable chunking above this memory threshold

Type:: float

progress_reporting

Enable progress reporting for large operations

Type:: bool

min_chunk_size

Minimum chunk size in data points

Type:: int

max_chunk_size

Maximum chunk size in data points

Type:: int

memory_limit_gb: float

gpu_memory_fraction: float | None

chunk_size_mb: int | None

out_of_memory_strategy: str

safety_factor: float

auto_chunk_threshold_gb: float

progress_reporting: bool

min_chunk_size: int

max_chunk_size: int

__post_init__()[source]: Validate configuration values.

__init__(memory_limit_gb=8.0, gpu_memory_fraction=None, chunk_size_mb=None, out_of_memory_strategy='fallback', safety_factor=0.8, auto_chunk_threshold_gb=4.0, progress_reporting=True, min_chunk_size=1000, max_chunk_size=1000000)

class nlsq.config.LargeDatasetConfig(enable_automatic_solver_selection=True, solver_selection_thresholds=<factory>)[source]

Bases: object

Configuration for large dataset processing.

enable_automatic_solver_selection

Automatically select optimal solver based on dataset size

Type:: bool

solver_selection_thresholds

Thresholds for automatic solver selection

Type:: Dict[str, int]

Notes

As of v0.2.0, all subsampling parameters have been removed. Use streaming optimization instead for unlimited datasets. See MIGRATION_V0.2.0.md for migration instructions.

enable_automatic_solver_selection: bool

solver_selection_thresholds: dict[str, int]

__post_init__()[source]

Validate configuration values.

__init__(enable_automatic_solver_selection=True, solver_selection_thresholds=<factory>)

class nlsq.config.JAXConfig[source]

Bases: object

Singleton configuration manager for JAX and memory settings.

This class ensures that JAX configuration is set once and consistently across all NLSQ modules, avoiding duplicate configuration calls. It also manages memory settings and large dataset configuration.

static __new__(cls)[source]

Ensure singleton pattern (thread-safe).

__init__()[source]: Initialize JAX and memory configuration if not already done.

classmethod enable_x64(enable=True)[source]

Enable or disable 64-bit precision.

Parameters:: enable (bool, optional) – If True, enable 64-bit precision. If False, use 32-bit. Default is True.

classmethod is_x64_enabled()[source]

Check if 64-bit precision is enabled.

Returns:: True if 64-bit precision is enabled, False otherwise.
Return type:: bool

classmethod precision_context(use_x64)[source]

Context manager for temporarily changing precision.

Parameters:: use_x64 (bool) – If True, use 64-bit precision within context. If False, use 32-bit precision.

Examples

>>> with JAXConfig.precision_context(use_x64=False):
...     # Code here runs with 32-bit precision
...     result = some_computation()
>>> # Back to previous precision setting

classmethod get_memory_config()[source]

Get the current memory configuration.

Returns:: Current memory configuration
Return type:: MemoryConfig

classmethod set_memory_config(config)[source]

Set the memory configuration.

Parameters:: config (MemoryConfig) – New memory configuration

classmethod get_large_dataset_config()[source]

Get the current large dataset configuration.

Returns:: Current large dataset configuration
Return type:: LargeDatasetConfig

classmethod set_large_dataset_config(config)[source]

Set the large dataset configuration.

Parameters:: config (LargeDatasetConfig) – New large dataset configuration

classmethod memory_context(memory_config)[source]

Context manager for temporarily changing memory configuration.

Parameters:: memory_config (MemoryConfig) – Temporary memory configuration

Examples

>>> from nlsq.config import JAXConfig, MemoryConfig
>>> temp_config = MemoryConfig(memory_limit_gb=16.0)
>>> with JAXConfig.memory_context(temp_config):
...     # Code here runs with increased memory limit
...     result = fit_large_dataset(func, x, y)
>>> # Back to previous memory settings

classmethod large_dataset_context(large_dataset_config)[source]

Context manager for temporarily changing large dataset configuration.

Parameters:: large_dataset_config (LargeDatasetConfig) – Temporary large dataset configuration

nlsq.config.enable_x64(enable=True)[source]

Enable or disable 64-bit precision.

Parameters:: enable (bool, optional) – If True, enable 64-bit precision. If False, use 32-bit. Default is True.

nlsq.config.is_x64_enabled()[source]

Check if 64-bit precision is enabled.

Returns:: True if 64-bit precision is enabled, False otherwise.
Return type:: bool

nlsq.config.precision_context(use_x64)[source]

Context manager for temporarily changing precision.

Parameters:: use_x64 (bool) – If True, use 64-bit precision within context. If False, use 32-bit precision.

Examples

>>> from nlsq.config import precision_context
>>> with precision_context(use_x64=False):
...     # Code here runs with 32-bit precision
...     result = some_computation()
>>> # Back to previous precision setting

nlsq.config.get_memory_config()[source]

Get the current memory configuration.

Returns:: Current memory configuration
Return type:: MemoryConfig

nlsq.config.set_memory_limits(memory_limit_gb, gpu_memory_fraction=None, safety_factor=0.8)[source]

Set memory limits for NLSQ operations.

Parameters:

memory_limit_gb (float) – Maximum memory to use in GB
gpu_memory_fraction (float, optional) – Fraction of GPU memory to use (0.0-1.0)
safety_factor (float, optional) – Safety factor for memory calculations (default: 0.8)

Examples

>>> from nlsq.config import set_memory_limits
>>> # Set 16GB memory limit with 80% GPU memory usage
>>> set_memory_limits(16.0, gpu_memory_fraction=0.8)

nlsq.config.configure_for_large_datasets(memory_limit_gb=8.0, enable_chunking=True, progress_reporting=True)[source]

Configure NLSQ for optimal large dataset performance.

This function sets up memory management, chunking, streaming, and other settings for handling large datasets efficiently.

Parameters:

memory_limit_gb (float, optional) – Maximum memory to use in GB (default: 8.0)
enable_chunking (bool, optional) – Enable automatic data chunking (default: True)
progress_reporting (bool, optional) – Enable progress reporting for long operations (default: True)

Notes

All large datasets use streaming optimization for 100% data utilization.

Examples

>>> from nlsq.config import configure_for_large_datasets
>>> # Configure for large datasets with 16GB memory limit
>>> configure_for_large_datasets(
...     memory_limit_gb=16.0,
...     progress_reporting=True
... )

nlsq.config.get_large_dataset_config()[source]

Get the current large dataset configuration.

Returns:: Current large dataset configuration
Return type:: LargeDatasetConfig

nlsq.config.memory_context(memory_config)[source]

Context manager for temporarily changing memory configuration.

Parameters:: memory_config (MemoryConfig) – Temporary memory configuration

Examples

>>> from nlsq.config import memory_context, MemoryConfig
>>> temp_config = MemoryConfig(memory_limit_gb=16.0)
>>> with memory_context(temp_config):
...     # Code here runs with increased memory limit
...     result = fit_large_dataset(func, x, y)
>>> # Back to previous memory settings

nlsq.config.large_dataset_context(large_dataset_config)[source]

Context manager for temporarily changing large dataset configuration.

Parameters:: large_dataset_config (LargeDatasetConfig) – Temporary large dataset configuration

Examples

>>> from nlsq.config import large_dataset_context, LargeDatasetConfig
>>> temp_config = LargeDatasetConfig(enable_automatic_solver_selection=True)
>>> with large_dataset_context(temp_config):
...     # Code here uses automatic solver selection
...     result = fit_large_dataset(func, x, y)

nlsq.config.get_jacobian_mode()[source]

Get Jacobian mode from configuration sources.

Configuration precedence (highest to lowest): 1. Environment variable (NLSQ_JACOBIAN_MODE) 2. Config file (~/.nlsq/config.json) 3. Auto-default

Returns:

mode (str) – Jacobian mode (‘auto’, ‘fwd’, or ‘rev’)
source (str) – Source of the configuration (‘environment variable’, ‘config file’, ‘auto-default’)

Return type:

tuple[str, str]

Examples

>>> from nlsq.config import get_jacobian_mode
>>> mode, source = get_jacobian_mode()
>>> print(f"Using {mode} mode from {source}")
Using auto mode from auto-default

Notes

Valid jacobian_mode values: - ‘auto’: Automatically select based on problem dimensions - ‘fwd’: Force forward-mode automatic differentiation (jacfwd) - ‘rev’: Force reverse-mode automatic differentiation (jacrev)

nlsq.config.set_jacobian_mode(mode)[source]

Set Jacobian mode via environment variable.

This sets the NLSQ_JACOBIAN_MODE environment variable for the current process. To persist the setting, use a config file at ~/.nlsq/config.json.

Parameters:: mode (str) – Jacobian mode (‘auto’, ‘fwd’, or ‘rev’)
Raises:: ValueError – If mode is not one of ‘auto’, ‘fwd’, ‘rev’

Examples

>>> from nlsq.config import set_jacobian_mode
>>> set_jacobian_mode('rev')  # Force reverse-mode AD for all fits