social-clusters/main.py

179 lines
5.1 KiB
Python
Raw Normal View History

2023-05-26 10:33:54 +04:00
#!/usr/bin/env python3
2023-05-29 22:56:53 +04:00
import json
2023-05-26 10:33:54 +04:00
import os
import sys
2023-05-29 22:56:53 +04:00
from datetime import date, datetime
2023-05-26 10:33:54 +04:00
2023-05-29 22:56:53 +04:00
import numpy as np
2023-05-26 10:33:54 +04:00
import pandas as pd
2023-05-29 22:56:53 +04:00
from geopy.extra.rate_limiter import RateLimiter
from geopy.geocoders import Nominatim
2023-05-26 10:33:54 +04:00
2023-05-29 22:56:53 +04:00
EMPTY_AGE = 0
UNIVERSITY_AGE = 21
SCHOOL_BEGIN_AGE = 7
SCHOOL_GRADUATED_AGE = 17
2023-05-26 10:33:54 +04:00
2023-05-29 22:56:53 +04:00
geolocator = Nominatim(user_agent="MyApp")
geocode = RateLimiter(geolocator.geocode, min_delay_seconds=1)
geo_cache = {}
def is_empty_str(value):
if value is None:
return True
return len(str(value).strip()) == 0
def is_empty_number(value):
if is_empty_str(value):
return True
str_val = str(value)
if str_val.startswith('-'):
str_val = str_val.replace('-', '', 1)
return not str_val.isnumeric()
def is_empty_collection(collection):
if is_empty_str(collection):
return True
if not isinstance(collection, list):
return True
return len(collection) == 0
def get_age(date_str):
if is_empty_str(date_str):
return EMPTY_AGE
today = date.today()
birthdate = datetime.strptime(date_str, '%d.%m.%Y')
age = today.year - birthdate.year - ((today.month, today.day) < (birthdate.month, birthdate.day))
return age
def get_years(year1, year2):
if year1 >= year2:
return year1 - year2
if year2 >= year1:
return year2 - year1
def get_age_from_education(education, value, additional_value):
if is_empty_collection(education):
return EMPTY_AGE
for item in education:
graduation = item[value]
if is_empty_number(graduation):
return EMPTY_AGE
return get_years(graduation, date.today().year) + additional_value
def prepare_dataset_age(df):
df['age'] = df.loc[:, 'bdate'].apply(get_age)
university_mask = (df['age'] == EMPTY_AGE) & (df['universities'].str.len() > 0)
df.loc[university_mask, 'age'] = df.loc[university_mask, 'universities'] \
.apply(lambda val: get_age_from_education(val, 'graduation', UNIVERSITY_AGE))
school_mask_1 = (df['age'] == EMPTY_AGE) & (df['schools'].str.len() > 0)
df.loc[school_mask_1, 'age'] = df.loc[school_mask_1, 'schools'] \
.apply(lambda val: get_age_from_education(val, 'year_graduated', SCHOOL_GRADUATED_AGE))
school_mask_2 = (df['age'] == EMPTY_AGE) & (df['schools'].str.len() > 0)
df.loc[school_mask_2, 'age'] = df.loc[school_mask_2, 'schools'] \
.apply(lambda val: get_age_from_education(val, 'year_from', SCHOOL_BEGIN_AGE))
return df
def prepare_dataset_status(df):
is_university_mask = ((df['age'] >= UNIVERSITY_AGE) | (df['age'] == EMPTY_AGE)) & \
((df['universities'].str.len() > 0) | (df['occupation_type'] == 'university'))
df['is_university'] = np.where(is_university_mask, True, False)
is_work_mask = ((df['age'] > SCHOOL_GRADUATED_AGE) | (df['age'] == EMPTY_AGE)) & \
((df['is_university']) | (df['occupation_type'] == 'work')) | \
(df['age'] > UNIVERSITY_AGE)
df['is_work'] = np.where(is_work_mask, True, False)
is_student_mask = ((df['occupation_type'] == 'university') &
((df['age'] >= SCHOOL_GRADUATED_AGE) & (df['age'] <= UNIVERSITY_AGE)))
df['is_student'] = np.where(is_student_mask, True, False)
is_schoolboy_mask = ((df['age'] < SCHOOL_GRADUATED_AGE) & (df['age'] != EMPTY_AGE)) | \
((df['age'] == EMPTY_AGE) & (df['occupation_type'] == 'school'))
df['is_schoolboy'] = np.where(is_schoolboy_mask, True, False)
return df
def load_geo_cache(json_file):
with open(json_file, 'r') as rf:
geo_cache.update(json.load(rf))
def save_geo_cache(json_file):
with open(json_file, 'w') as wf:
json.dump(geo_cache, wf)
print('Geocache saved')
def update_geo_cache(cities, json_file):
is_changed = False
for city in cities:
if is_empty_str(city):
continue
result = geo_cache.get(city)
if result is not None:
continue
print(f'{len(geo_cache.keys())}/{len(cities)} - Try to load geocode for {city}')
location = geocode(city)
result = (location.latitude, location.longitude)
geo_cache[city] = result
is_changed = True
if len(geo_cache.keys()) % 50 == 0:
save_geo_cache(json_file)
if is_changed:
save_geo_cache(json_file)
def prepare_dataset_location(df):
json_file = 'geocache.json'
load_geo_cache(json_file)
update_geo_cache(df['city'].unique().tolist(), json_file)
df['location'] = df['city'] \
.apply(lambda val: '' if is_empty_str(val) else geo_cache[val])
return df
def prepare_dataset(json_file):
2023-05-26 10:33:54 +04:00
df = pd.read_json(json_file)
2023-05-29 22:56:53 +04:00
df = prepare_dataset_age(df)
df = prepare_dataset_status(df)
df = prepare_dataset_location(df)
return df
def __main(json_file):
df = prepare_dataset(json_file)
2023-05-26 10:33:54 +04:00
print('done')
if __name__ == '__main__':
if len(sys.argv) != 2:
print('You must specify the raw_dataset json file')
exit(1)
if not os.path.isfile(sys.argv[1]):
print(f'File {sys.argv[1]} is not exists')
__main(sys.argv[1])